<commit_before>/*
 * Copyright (c) 2013, 2014 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"github.com/conformal/btcjson"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
)

// RescanMsg is the interface type for messages sent to the
// RescanManager's message channel.
type RescanMsg interface {
	ImplementsRescanMsg()
}

// RescanStartedMsg reports the job being processed for a new
// rescan.
type RescanStartedMsg RescanJob

// ImplementsRescanMsg is implemented to satisify the RescanMsg
// interface.
func (r *RescanStartedMsg) ImplementsRescanMsg() {}

// RescanProgressMsg reports the current progress made by a rescan
// for a set of account's addresses.
type RescanProgressMsg struct {
	Addresses map[*Account][]btcutil.Address
	Height    int32
}

// ImplementsRescanMsg is implemented to satisify the RescanMsg
// interface.
func (r *RescanProgressMsg) ImplementsRescanMsg() {}

// RescanFinishedMsg reports the set of account's addresses of a
// possibly-finished rescan, or an error if the rescan failed.
type RescanFinishedMsg struct {
	Addresses map[*Account][]btcutil.Address
	Error     *btcjson.Error
}

// ImplementsRescanMsg is implemented to satisify the RescanMsg
// interface.
func (r *RescanFinishedMsg) ImplementsRescanMsg() {}

// RescanManager manages a set of current and to be processed account's
// addresses, batching waiting jobs together to minimize the total time
// needed to rescan many separate jobs.  Rescan requests are processed
// one at a time, and the next batch does not run until the current
// has finished.
type RescanManager struct {
	addJob          chan *RescanJob
	sendJob         chan *RescanJob
	status          chan interface{} // rescanProgress and rescanFinished
	msgs            chan RescanMsg
	jobCompleteChan chan chan struct{}
}

// NewRescanManager creates a new RescanManger.  If msgChan is non-nil,
// rescan messages are sent to the channel for additional processing by
// the caller.
func NewRescanManager(msgChan chan RescanMsg) *RescanManager {
	return &RescanManager{
		addJob:          make(chan *RescanJob, 1),
		sendJob:         make(chan *RescanJob, 1),
		status:          make(chan interface{}, 1),
		msgs:            msgChan,
		jobCompleteChan: make(chan chan struct{}, 1),
	}
}

// Start starts the goroutines to run the RescanManager.
func (m *RescanManager) Start() {
	go m.jobHandler()
	go m.rpcHandler()
}

type rescanBatch struct {
	addrs     map[*Account][]btcutil.Address
	outpoints map[btcwire.OutPoint]struct{}
	height    int32
	complete  chan struct{}
}

func newRescanBatch() *rescanBatch {
	return &rescanBatch{
		addrs:     map[*Account][]btcutil.Address{},
		outpoints: map[btcwire.OutPoint]struct{}{},
		height:    -1,
		complete:  make(chan struct{}),
	}
}

func (b *rescanBatch) done() {
	close(b.complete)
}

func (b *rescanBatch) empty() bool {
	return len(b.addrs) == 0
}

func (b *rescanBatch) job() *RescanJob {
	// Create slice of outpoint points from the batch's set.
	outpoints := make([]*btcwire.OutPoint, 0, len(b.outpoints))
	for outpoint := range b.outpoints {
		opCopy := outpoint
		outpoints = append(outpoints, &opCopy)
	}

	return &RescanJob{
		Addresses:   b.addrs,
		OutPoints:   outpoints,
		StartHeight: b.height,
	}
}

func (b *rescanBatch) merge(job *RescanJob) {
	for acct, addr := range job.Addresses {
		b.addrs[acct] = append(b.addrs[acct], addr...)
	}
	for _, op := range job.OutPoints {
		b.outpoints[*op] = struct{}{}
	}
	if b.height == -1 || job.StartHeight < b.height {
		b.height = job.StartHeight
	}
}

// Status types for the handler.
type rescanProgress int32
type rescanFinished *btcjson.Error

// jobHandler runs the RescanManager's for-select loop to manage rescan jobs
// and dispatch requests.
func (m *RescanManager) jobHandler() {
	curBatch := newRescanBatch()
	nextBatch := newRescanBatch()

	for {
		select {
		case job := <-m.addJob:
			if curBatch.empty() {
				// Set current batch as this job and send
				// request.
				curBatch.merge(job)
				m.sendJob <- job

				// Send the channel that is closed when the
				// current batch completes.
				m.jobCompleteChan <- curBatch.complete

				// Notify listener of a newly-started rescan.
				if m.msgs != nil {
					m.msgs <- (*RescanStartedMsg)(job)
				}
			} else {
				// Add job to waiting batch.
				nextBatch.merge(job)

				// Send the channel that is closed when the
				// waiting batch completes.
				m.jobCompleteChan <- nextBatch.complete
			}

		case status := <-m.status:
			switch s := status.(type) {
			case rescanProgress:
				if m.msgs != nil {
					m.msgs <- &RescanProgressMsg{
						Addresses: curBatch.addrs,
						Height:    int32(s),
					}
				}

			case rescanFinished:
				if m.msgs != nil {
					m.msgs <- &RescanFinishedMsg{
						Addresses: curBatch.addrs,
						Error:     (*btcjson.Error)(s),
					}
				}
				curBatch.done()

				curBatch, nextBatch = nextBatch, newRescanBatch()

				if !curBatch.empty() {
					job := curBatch.job()
					m.sendJob <- job
					if m.msgs != nil {
						m.msgs <- (*RescanStartedMsg)(job)
					}
				}
			}
		}
	}
}

// rpcHandler reads jobs sent by the jobHandler and sends the rpc requests
// to perform the rescan.  New jobs are not read until a rescan finishes.
// The jobHandler is notified when the processing the rescan finishes.
func (m *RescanManager) rpcHandler() {
	for job := range m.sendJob {
		var addrStrs []string
		for _, addrs := range job.Addresses {
			for i := range addrs {
				addrStrs = append(addrStrs, addrs[i].EncodeAddress())
			}
		}

		c := CurrentServerConn()
		jsonErr := Rescan(c, job.StartHeight, addrStrs, job.OutPoints)
		m.status <- rescanFinished(jsonErr)
	}
}

// RescanJob is a job to be processed by the RescanManager.  The job includes
// a set of account's addresses, a starting height to begin the rescan, and
// outpoints spendable by the addresses thought to be unspent.
type RescanJob struct {
	Addresses   map[*Account][]btcutil.Address
	OutPoints   []*btcwire.OutPoint
	StartHeight int32
}

// Merge merges the work from k into j, setting the starting height to
// the minimum of the two jobs.  This method does not check for
// duplicate addresses or outpoints.
func (j *RescanJob) Merge(k *RescanJob) {
	for acct, addrs := range k.Addresses {
		j.Addresses[acct] = append(j.Addresses[acct], addrs...)
	}
	for _, op := range k.OutPoints {
		j.OutPoints = append(j.OutPoints, op)
	}
	if k.StartHeight < j.StartHeight {
		j.StartHeight = k.StartHeight
	}
}

// SubmitJob submits a RescanJob to the RescanManager.  A channel is returned
// that is closed once the rescan request for the job completes.
func (m *RescanManager) SubmitJob(job *RescanJob) <-chan struct{} {
	m.addJob <- job
	return <-m.jobCompleteChan
}

// MarkProgress messages the RescanManager with the height of the block
// last processed by a running rescan.
func (m *RescanManager) MarkProgress(height int32) {
	m.status <- rescanProgress(height)
}
<commit_msg>Make a comment understandable.<commit_after>/*
 * Copyright (c) 2013, 2014 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"github.com/conformal/btcjson"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
)

// RescanMsg is the interface type for messages sent to the
// RescanManager's message channel.
type RescanMsg interface {
	ImplementsRescanMsg()
}

// RescanStartedMsg reports the job being processed for a new
// rescan.
type RescanStartedMsg RescanJob

// ImplementsRescanMsg is implemented to satisify the RescanMsg
// interface.
func (r *RescanStartedMsg) ImplementsRescanMsg() {}

// RescanProgressMsg reports the current progress made by a rescan
// for a set of account's addresses.
type RescanProgressMsg struct {
	Addresses map[*Account][]btcutil.Address
	Height    int32
}

// ImplementsRescanMsg is implemented to satisify the RescanMsg
// interface.
func (r *RescanProgressMsg) ImplementsRescanMsg() {}

// RescanFinishedMsg reports the set of account's addresses of a
// possibly-finished rescan, or an error if the rescan failed.
type RescanFinishedMsg struct {
	Addresses map[*Account][]btcutil.Address
	Error     *btcjson.Error
}

// ImplementsRescanMsg is implemented to satisify the RescanMsg
// interface.
func (r *RescanFinishedMsg) ImplementsRescanMsg() {}

// RescanManager manages a set of current and to be processed account's
// addresses, batching waiting jobs together to minimize the total time
// needed to rescan many separate jobs.  Rescan requests are processed
// one at a time, and the next batch does not run until the current
// has finished.
type RescanManager struct {
	addJob          chan *RescanJob
	sendJob         chan *RescanJob
	status          chan interface{} // rescanProgress and rescanFinished
	msgs            chan RescanMsg
	jobCompleteChan chan chan struct{}
}

// NewRescanManager creates a new RescanManger.  If msgChan is non-nil,
// rescan messages are sent to the channel for additional processing by
// the caller.
func NewRescanManager(msgChan chan RescanMsg) *RescanManager {
	return &RescanManager{
		addJob:          make(chan *RescanJob, 1),
		sendJob:         make(chan *RescanJob, 1),
		status:          make(chan interface{}, 1),
		msgs:            msgChan,
		jobCompleteChan: make(chan chan struct{}, 1),
	}
}

// Start starts the goroutines to run the RescanManager.
func (m *RescanManager) Start() {
	go m.jobHandler()
	go m.rpcHandler()
}

type rescanBatch struct {
	addrs     map[*Account][]btcutil.Address
	outpoints map[btcwire.OutPoint]struct{}
	height    int32
	complete  chan struct{}
}

func newRescanBatch() *rescanBatch {
	return &rescanBatch{
		addrs:     map[*Account][]btcutil.Address{},
		outpoints: map[btcwire.OutPoint]struct{}{},
		height:    -1,
		complete:  make(chan struct{}),
	}
}

func (b *rescanBatch) done() {
	close(b.complete)
}

func (b *rescanBatch) empty() bool {
	return len(b.addrs) == 0
}

func (b *rescanBatch) job() *RescanJob {
	// Create slice of outpoints from the batch's set.
	outpoints := make([]*btcwire.OutPoint, 0, len(b.outpoints))
	for outpoint := range b.outpoints {
		opCopy := outpoint
		outpoints = append(outpoints, &opCopy)
	}

	return &RescanJob{
		Addresses:   b.addrs,
		OutPoints:   outpoints,
		StartHeight: b.height,
	}
}

func (b *rescanBatch) merge(job *RescanJob) {
	for acct, addr := range job.Addresses {
		b.addrs[acct] = append(b.addrs[acct], addr...)
	}
	for _, op := range job.OutPoints {
		b.outpoints[*op] = struct{}{}
	}
	if b.height == -1 || job.StartHeight < b.height {
		b.height = job.StartHeight
	}
}

// Status types for the handler.
type rescanProgress int32
type rescanFinished *btcjson.Error

// jobHandler runs the RescanManager's for-select loop to manage rescan jobs
// and dispatch requests.
func (m *RescanManager) jobHandler() {
	curBatch := newRescanBatch()
	nextBatch := newRescanBatch()

	for {
		select {
		case job := <-m.addJob:
			if curBatch.empty() {
				// Set current batch as this job and send
				// request.
				curBatch.merge(job)
				m.sendJob <- job

				// Send the channel that is closed when the
				// current batch completes.
				m.jobCompleteChan <- curBatch.complete

				// Notify listener of a newly-started rescan.
				if m.msgs != nil {
					m.msgs <- (*RescanStartedMsg)(job)
				}
			} else {
				// Add job to waiting batch.
				nextBatch.merge(job)

				// Send the channel that is closed when the
				// waiting batch completes.
				m.jobCompleteChan <- nextBatch.complete
			}

		case status := <-m.status:
			switch s := status.(type) {
			case rescanProgress:
				if m.msgs != nil {
					m.msgs <- &RescanProgressMsg{
						Addresses: curBatch.addrs,
						Height:    int32(s),
					}
				}

			case rescanFinished:
				if m.msgs != nil {
					m.msgs <- &RescanFinishedMsg{
						Addresses: curBatch.addrs,
						Error:     (*btcjson.Error)(s),
					}
				}
				curBatch.done()

				curBatch, nextBatch = nextBatch, newRescanBatch()

				if !curBatch.empty() {
					job := curBatch.job()
					m.sendJob <- job
					if m.msgs != nil {
						m.msgs <- (*RescanStartedMsg)(job)
					}
				}
			}
		}
	}
}

// rpcHandler reads jobs sent by the jobHandler and sends the rpc requests
// to perform the rescan.  New jobs are not read until a rescan finishes.
// The jobHandler is notified when the processing the rescan finishes.
func (m *RescanManager) rpcHandler() {
	for job := range m.sendJob {
		var addrStrs []string
		for _, addrs := range job.Addresses {
			for i := range addrs {
				addrStrs = append(addrStrs, addrs[i].EncodeAddress())
			}
		}

		c := CurrentServerConn()
		jsonErr := Rescan(c, job.StartHeight, addrStrs, job.OutPoints)
		m.status <- rescanFinished(jsonErr)
	}
}

// RescanJob is a job to be processed by the RescanManager.  The job includes
// a set of account's addresses, a starting height to begin the rescan, and
// outpoints spendable by the addresses thought to be unspent.
type RescanJob struct {
	Addresses   map[*Account][]btcutil.Address
	OutPoints   []*btcwire.OutPoint
	StartHeight int32
}

// Merge merges the work from k into j, setting the starting height to
// the minimum of the two jobs.  This method does not check for
// duplicate addresses or outpoints.
func (j *RescanJob) Merge(k *RescanJob) {
	for acct, addrs := range k.Addresses {
		j.Addresses[acct] = append(j.Addresses[acct], addrs...)
	}
	for _, op := range k.OutPoints {
		j.OutPoints = append(j.OutPoints, op)
	}
	if k.StartHeight < j.StartHeight {
		j.StartHeight = k.StartHeight
	}
}

// SubmitJob submits a RescanJob to the RescanManager.  A channel is returned
// that is closed once the rescan request for the job completes.
func (m *RescanManager) SubmitJob(job *RescanJob) <-chan struct{} {
	m.addJob <- job
	return <-m.jobCompleteChan
}

// MarkProgress messages the RescanManager with the height of the block
// last processed by a running rescan.
func (m *RescanManager) MarkProgress(height int32) {
	m.status <- rescanProgress(height)
}
<|endoftext|>


<commit_before>package staert

import (
	"github.com/BurntSushi/toml"
	"github.com/containous/flaeg"
	"os"
	"path/filepath"
	"strings"
)

// Source interface must be satisfy to Add any kink of Source to Staert as like as TomlFile or Flaeg
type Source interface {
	Parse(cmd *flaeg.Command) (*flaeg.Command, error)
}

// Staert contains the struct to configure, thee default values inside structs and the sources
type Staert struct {
	command *flaeg.Command
	sources []Source
}

// NewStaert creats and return a pointer on Staert. Need defaultConfig and defaultPointersConfig given by references
func NewStaert(rootCommand *flaeg.Command) *Staert {
	s := Staert{
		command: rootCommand,
	}
	return &s
}

// AddSource adds new Source to Staert, give it by reference
func (s *Staert) AddSource(src Source) {
	s.sources = append(s.sources, src)
}

// getConfig for a flaeg.Command run sources Parse func in the raw
func (s *Staert) parseConfigAllSources(cmd *flaeg.Command) error {
	for _, src := range s.sources {
		var err error
		_, err = src.Parse(cmd)
		if err != nil {
			return err
		}
	}
	return nil
}

// LoadConfig check which command is called and parses config
// It returns the the parsed config or an error if it fails
func (s *Staert) LoadConfig() (interface{}, error) {
	for _, src := range s.sources {
		//Type assertion
		f, ok := src.(*flaeg.Flaeg)
		if ok {
			if fCmd, err := f.GetCommand(); err != nil {
				return nil, err
			} else if s.command != fCmd {
				//IF fleag sub-command
				s.command, err = f.Parse(fCmd)
				return s.command.Config, err
			}
		}
	}
	err := s.parseConfigAllSources(s.command)
	return s.command.Config, err
}

// Run calls the Run func of the command
// Warning, Run doesn't parse the config
func (s *Staert) Run() error {
	return s.command.Run()
}

//TomlSource impement Source
type TomlSource struct {
	filename     string
	dirNfullpath []string
	fullpath     string
}

// NewTomlSource creats and return a pointer on TomlSource.
// Parameter filename is the file name (without extension type, ".toml" will be added)
// dirNfullpath may contain directories or fullpath to the file.
func NewTomlSource(filename string, dirNfullpath []string) *TomlSource {
	return &TomlSource{filename, dirNfullpath, ""}
}

// ConfigFileUsed return config file used
func (ts *TomlSource) ConfigFileUsed() string {
	return ts.fullpath
}

func preprocessDir(dirIn string) (string, error) {
	dirOut := dirIn
	if strings.HasPrefix(dirIn, "$") {
		end := strings.Index(dirIn, string(os.PathSeparator))
		if end == -1 {
			end = len(dirIn)
		}
		dirOut = os.Getenv(dirIn[1:end]) + dirIn[end:]
	}
	dirOut, err := filepath.Abs(dirOut)
	return dirOut, err
}

func findFile(filename string, dirNfile []string) string {
	for _, df := range dirNfile {
		if df != "" {
			fullpath, _ := preprocessDir(df)
			if fileinfo, err := os.Stat(fullpath); err == nil && !fileinfo.IsDir() {
				return fullpath
			}
			fullpath = fullpath + "/" + filename + ".toml"
			if fileinfo, err := os.Stat(fullpath); err == nil && !fileinfo.IsDir() {
				return fullpath
			}
		}
	}
	return ""
}

// Parse calls Flaeg Load Function
func (ts *TomlSource) Parse(cmd *flaeg.Command) (*flaeg.Command, error) {
	ts.fullpath = findFile(ts.filename, ts.dirNfullpath)
	if len(ts.fullpath) < 2 {
		return cmd, nil
	}
	metadata, err := toml.DecodeFile(ts.fullpath, cmd.Config)
	if err != nil {
		return nil, err
	}
	flaegArgs := []string{}
	keys := metadata.Keys()
	for i, key := range keys {
		// fmt.Println(key)
		if metadata.Type(key.String()) == "Hash" {
			//Ptr case
			// fmt.Printf("%s is a ptr\n", key)
			hasUnderField := false
			for j := i; j < len(keys); j++ {
				// fmt.Printf("%s =? %s\n", keys[j].String(), "."+key.String())
				if strings.Contains(keys[j].String(), key.String()+".") {
					hasUnderField = true
					break
				}
			}
			if !hasUnderField {
				flaegArgs = append(flaegArgs, "--"+strings.ToLower(key.String()))
			}
		}
	}
	// fmt.Println(flaegArgs)
	err = flaeg.Load(cmd.Config, cmd.DefaultPointersConfig, flaegArgs)
	//if err!= missing parser err
	if err != nil && !strings.Contains(err.Error(), "No parser for type") {
		return nil, err
	}
	return cmd, nil
}
<commit_msg>fix toml Parse func (#12)<commit_after>package staert

import (
	"github.com/BurntSushi/toml"
	"github.com/containous/flaeg"
	"os"
	"path/filepath"
	"strings"
)

// Source interface must be satisfy to Add any kink of Source to Staert as like as TomlFile or Flaeg
type Source interface {
	Parse(cmd *flaeg.Command) (*flaeg.Command, error)
}

// Staert contains the struct to configure, thee default values inside structs and the sources
type Staert struct {
	command *flaeg.Command
	sources []Source
}

// NewStaert creats and return a pointer on Staert. Need defaultConfig and defaultPointersConfig given by references
func NewStaert(rootCommand *flaeg.Command) *Staert {
	s := Staert{
		command: rootCommand,
	}
	return &s
}

// AddSource adds new Source to Staert, give it by reference
func (s *Staert) AddSource(src Source) {
	s.sources = append(s.sources, src)
}

// getConfig for a flaeg.Command run sources Parse func in the raw
func (s *Staert) parseConfigAllSources(cmd *flaeg.Command) error {
	for _, src := range s.sources {
		var err error
		_, err = src.Parse(cmd)
		if err != nil {
			return err
		}
	}
	return nil
}

// LoadConfig check which command is called and parses config
// It returns the the parsed config or an error if it fails
func (s *Staert) LoadConfig() (interface{}, error) {
	for _, src := range s.sources {
		//Type assertion
		f, ok := src.(*flaeg.Flaeg)
		if ok {
			if fCmd, err := f.GetCommand(); err != nil {
				return nil, err
			} else if s.command != fCmd {
				//IF fleag sub-command
				s.command, err = f.Parse(fCmd)
				return s.command.Config, err
			}
		}
	}
	err := s.parseConfigAllSources(s.command)
	return s.command.Config, err
}

// Run calls the Run func of the command
// Warning, Run doesn't parse the config
func (s *Staert) Run() error {
	return s.command.Run()
}

//TomlSource impement Source
type TomlSource struct {
	filename     string
	dirNfullpath []string
	fullpath     string
}

// NewTomlSource creats and return a pointer on TomlSource.
// Parameter filename is the file name (without extension type, ".toml" will be added)
// dirNfullpath may contain directories or fullpath to the file.
func NewTomlSource(filename string, dirNfullpath []string) *TomlSource {
	return &TomlSource{filename, dirNfullpath, ""}
}

// ConfigFileUsed return config file used
func (ts *TomlSource) ConfigFileUsed() string {
	return ts.fullpath
}

func preprocessDir(dirIn string) (string, error) {
	dirOut := dirIn
	if strings.HasPrefix(dirIn, "$") {
		end := strings.Index(dirIn, string(os.PathSeparator))
		if end == -1 {
			end = len(dirIn)
		}
		dirOut = os.Getenv(dirIn[1:end]) + dirIn[end:]
	}
	dirOut, err := filepath.Abs(dirOut)
	return dirOut, err
}

func findFile(filename string, dirNfile []string) string {
	for _, df := range dirNfile {
		if df != "" {
			fullpath, _ := preprocessDir(df)
			if fileinfo, err := os.Stat(fullpath); err == nil && !fileinfo.IsDir() {
				return fullpath
			}
			fullpath = fullpath + "/" + filename + ".toml"
			if fileinfo, err := os.Stat(fullpath); err == nil && !fileinfo.IsDir() {
				return fullpath
			}
		}
	}
	return ""
}

// Parse calls Flaeg Load Function
func (ts *TomlSource) Parse(cmd *flaeg.Command) (*flaeg.Command, error) {
	ts.fullpath = findFile(ts.filename, ts.dirNfullpath)
	if len(ts.fullpath) < 2 {
		return cmd, nil
	}
	metadata, err := toml.DecodeFile(ts.fullpath, cmd.Config)
	if err != nil {
		return nil, err
	}
	flags, err := flaeg.GetFlags(cmd.Config)
	if err != nil {
		return nil, err
	}
	flaegArgs, err := generateArgs(metadata, flags)
	if err != nil {
		return nil, err
	}

	// fmt.Println(flaegArgs)
	err = flaeg.Load(cmd.Config, cmd.DefaultPointersConfig, flaegArgs)
	//if err!= missing parser err
	if err != nil && !strings.Contains(err.Error(), "No parser for type") {
		return nil, err
	}
	return cmd, nil
}

func generateArgs(metadata toml.MetaData, flags []string) ([]string, error) {
	flaegArgs := []string{}
	keys := metadata.Keys()
	for i, key := range keys {
		// fmt.Println(key)
		if metadata.Type(key.String()) == "Hash" {
			//Ptr case
			// fmt.Printf("%s is a ptr\n", key)
			hasUnderField := false
			for j := i; j < len(keys); j++ {
				// fmt.Printf("%s =? %s\n", keys[j].String(), "."+key.String())
				if strings.Contains(keys[j].String(), key.String()+".") {
					hasUnderField = true
					break
				}
			}
			if !hasUnderField {
				match := false
				for _, flag := range flags {
					if flag == strings.ToLower(key.String()) {
						match = true
						break
					}
				}
				if match {
					flaegArgs = append(flaegArgs, "--"+strings.ToLower(key.String()))
				}
			}
		}
	}
	return flaegArgs, nil
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/asim/mq/broker"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func pub(w http.ResponseWriter, r *http.Request) {
	topic := r.URL.Query().Get("topic")

	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Pub error", http.StatusInternalServerError)
		return
	}
	r.Body.Close()

	if err := broker.Publish(topic, b); err != nil {
		http.Error(w, "Pub error", http.StatusInternalServerError)
	}
}

func sub(w http.ResponseWriter, r *http.Request) {
	var wr writer

	if websocket.IsWebSocketUpgrade(r) {
		conn, err := upgrader.Upgrade(w, r, w.Header())
		if err != nil {
			return
		}
		wr = &wsWriter{conn}
	} else {
		wr = &httpWriter{w}
	}

	topic := r.URL.Query().Get("topic")

	ch, err := broker.Subscribe(topic)
	if err != nil {
		http.Error(w, fmt.Sprintf("Could not retrieve events: %v", err), http.StatusInternalServerError)
		return
	}
	defer broker.Unsubscribe(topic, ch)

	for {
		select {
		case e := <-ch:
			if err = wr.Write(e); err != nil {
				return
			}
		}
	}
}
<commit_msg>Support for websocket publish<commit_after>package http

import (
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/asim/mq/broker"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func pub(w http.ResponseWriter, r *http.Request) {
	topic := r.URL.Query().Get("topic")

	if websocket.IsWebSocketUpgrade(r) {
		conn, err := upgrader.Upgrade(w, r, w.Header())
		if err != nil {
			return
		}
		for {
			messageType, b, err := conn.ReadMessage()
			if messageType == -1 {
				return
			}
			if err != nil {
				continue
			}
			broker.Publish(topic, b)
		}
	} else {
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, "Pub error", http.StatusInternalServerError)
			return
		}
		r.Body.Close()
		if err := broker.Publish(topic, b); err != nil {
			http.Error(w, "Pub error", http.StatusInternalServerError)
		}
	}
}

func sub(w http.ResponseWriter, r *http.Request) {
	var wr writer

	if websocket.IsWebSocketUpgrade(r) {
		conn, err := upgrader.Upgrade(w, r, w.Header())
		if err != nil {
			return
		}
		wr = &wsWriter{conn}
	} else {
		wr = &httpWriter{w}
	}

	topic := r.URL.Query().Get("topic")

	ch, err := broker.Subscribe(topic)
	if err != nil {
		http.Error(w, fmt.Sprintf("Could not retrieve events: %v", err), http.StatusInternalServerError)
		return
	}
	defer broker.Unsubscribe(topic, ch)

	for {
		select {
		case e := <-ch:
			if err = wr.Write(e); err != nil {
				return
			}
		}
	}
}
<|endoftext|>


<commit_before>/* Copyright (c) 2017 Jason Ish
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. 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 ``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 AUTHOR 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 evereader

import (
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/eve"
	"github.com/jasonish/evebox/log"
	"io"
	"sync"
	"time"
)

const BATCH_SIZE = 1000

// EveFileProcessor processes eve files by reading events from reader, applying
// any filters then sending the events to an event sink.
type EveFileProcessor struct {
	Filename          string
	BookmarkDirectory string
	Sink              core.EveEventSink

	filters []eve.EveFilter

	customFields map[string]interface{}

	bookmarker *Bookmarker

	// Total number of events processed.
	count uint64

	// Flag to set to stop processing.
	stop bool

	wg sync.WaitGroup

	// Internal metrics counting.
	lastStatCount uint64
	lastStatTime  time.Time
	eofs          uint64
}

func (p *EveFileProcessor) AddFilter(filter eve.EveFilter) {
	p.filters = append(p.filters, filter)
}

func (p *EveFileProcessor) AddCustomField(field string, value interface{}) {
	if p.customFields == nil {
		p.customFields = make(map[string]interface{})
	}
	p.customFields[field] = value
}

func (p *EveFileProcessor) Start() {
	p.lastStatTime = time.Now()
	p.wg.Add(1)
	go p.Run()
}

func (p *EveFileProcessor) Stop() {
	p.stop = true
	p.wg.Wait()
}

func (p *EveFileProcessor) Run() {
	// Outer loop retries opening the file in case of failures...
	for {
		var bookmarker *Bookmarker

		reader, err := NewFollowingReader(p.Filename)
		if err != nil {
			log.Warning("Failed to open %s (will try again): %v",
				p.Filename, err)
			goto Retry
		}

		bookmarker, err = NewBookmarker(reader, p.BookmarkDirectory)
		if err != nil {
			log.Warning("Failed to get bookmarker (will try again): %v", err)
			reader.Close()
			goto Retry
		}

		if err := p.process(reader, bookmarker); err != nil {
			log.Error("Processing error, will retry: %v", err)
			goto Retry
		}

	Retry:
		if p.stop {
			break
		}
		time.Sleep(1 * time.Second)
	}

	p.wg.Done()
}

func (p *EveFileProcessor) process(reader *FollowingReader, bookmarker *Bookmarker) error {

	count := uint64(0)

	for {
		eof := false

		event, err := reader.Next()
		if err != nil {
			if err == io.EOF {
				eof = true
				p.eofs++
			} else {
				return err
			}
		}

		if !eof {
			for _, filter := range p.filters {
				filter.Filter(event)
			}
			p.addCustomFields(event)
			if err := p.Sink.Submit(event); err != nil {
				log.Error("Failed to submit event: %v", err)
				continue
			}
			count++
		}

		// On every EOF or batch size, commit.
		if (eof && count > 0) || (count > 0 && count%BATCH_SIZE == 0) {
			start := time.Now()
			if err := p.commit(); err != nil {
				log.Error("Commit failed: %v", err)
				return err
			}
			log.Debug("Committed %d events in %v", count,
				time.Now().Sub(start))
			bookmarker.UpdateBookmark()
			p.count += count
			count = 0
		}

		// Print stats.
		now := time.Now()
		if now.Sub(p.lastStatTime).Seconds() > 60 {
			log.Info("Total: %d; last minute: %d; EOFs: %d",
				p.count,
				p.count-p.lastStatCount,
				p.eofs)
			p.lastStatCount = count
			p.lastStatTime = now
			p.eofs = 0
		}

		if p.stop {
			break
		}

		// If eof, sleep for a moment.
		if eof {
			time.Sleep(1 * time.Second)
		}
	}

	return nil
}

func (p *EveFileProcessor) commit() error {
	for {
		_, err := p.Sink.Commit()
		if err == nil {
			return nil
		}
		if p.stop {
			return err
		}
		log.Error("Failed to commit events, will try again: %v", err)
		time.Sleep(1 * time.Second)
		continue
	}
}

func (p *EveFileProcessor) addCustomFields(event eve.EveEvent) {
	for key := range p.customFields {
		event[key] = p.customFields[key]
	}
}
<commit_msg>eve reader: fix per minute stats count<commit_after>/* Copyright (c) 2017 Jason Ish
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. 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 ``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 AUTHOR 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 evereader

import (
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/eve"
	"github.com/jasonish/evebox/log"
	"io"
	"sync"
	"time"
)

const BATCH_SIZE = 1000

// EveFileProcessor processes eve files by reading events from reader, applying
// any filters then sending the events to an event sink.
type EveFileProcessor struct {
	Filename          string
	BookmarkDirectory string
	Sink              core.EveEventSink

	filters []eve.EveFilter

	customFields map[string]interface{}

	bookmarker *Bookmarker

	// Total number of events processed.
	count uint64

	// Flag to set to stop processing.
	stop bool

	wg sync.WaitGroup

	// Internal metrics counting.
	lastStatCount uint64
	lastStatTime  time.Time
	eofs          uint64
}

func (p *EveFileProcessor) AddFilter(filter eve.EveFilter) {
	p.filters = append(p.filters, filter)
}

func (p *EveFileProcessor) AddCustomField(field string, value interface{}) {
	if p.customFields == nil {
		p.customFields = make(map[string]interface{})
	}
	p.customFields[field] = value
}

func (p *EveFileProcessor) Start() {
	p.lastStatTime = time.Now()
	p.wg.Add(1)
	go p.Run()
}

func (p *EveFileProcessor) Stop() {
	p.stop = true
	p.wg.Wait()
}

func (p *EveFileProcessor) Run() {
	// Outer loop retries opening the file in case of failures...
	for {
		var bookmarker *Bookmarker

		reader, err := NewFollowingReader(p.Filename)
		if err != nil {
			log.Warning("Failed to open %s (will try again): %v",
				p.Filename, err)
			goto Retry
		}

		bookmarker, err = NewBookmarker(reader, p.BookmarkDirectory)
		if err != nil {
			log.Warning("Failed to get bookmarker (will try again): %v", err)
			reader.Close()
			goto Retry
		}

		if err := p.process(reader, bookmarker); err != nil {
			log.Error("Processing error, will retry: %v", err)
			goto Retry
		}

	Retry:
		if p.stop {
			break
		}
		time.Sleep(1 * time.Second)
	}

	p.wg.Done()
}

func (p *EveFileProcessor) process(reader *FollowingReader, bookmarker *Bookmarker) error {

	count := uint64(0)

	for {
		eof := false

		event, err := reader.Next()
		if err != nil {
			if err == io.EOF {
				eof = true
				p.eofs++
			} else {
				return err
			}
		}

		if !eof {
			for _, filter := range p.filters {
				filter.Filter(event)
			}
			p.addCustomFields(event)
			if err := p.Sink.Submit(event); err != nil {
				log.Error("Failed to submit event: %v", err)
				continue
			}
			count++
		}

		// On every EOF or batch size, commit.
		if (eof && count > 0) || (count > 0 && count%BATCH_SIZE == 0) {
			start := time.Now()
			if err := p.commit(); err != nil {
				log.Error("Commit failed: %v", err)
				return err
			}
			log.Debug("Committed %d events in %v", count,
				time.Now().Sub(start))
			bookmarker.UpdateBookmark()
			p.count += count
			count = 0
		}

		// Print stats.
		now := time.Now()
		if now.Sub(p.lastStatTime).Seconds() > 60 {
			log.Info("Total: %d; last minute: %d; EOFs: %d",
				p.count,
				p.count-p.lastStatCount,
				p.eofs)
			p.lastStatCount = p.count
			p.lastStatTime = now
			p.eofs = 0
		}

		if p.stop {
			break
		}

		// If eof, sleep for a moment.
		if eof {
			time.Sleep(1 * time.Second)
		}
	}

	return nil
}

func (p *EveFileProcessor) commit() error {
	for {
		_, err := p.Sink.Commit()
		if err == nil {
			return nil
		}
		if p.stop {
			return err
		}
		log.Error("Failed to commit events, will try again: %v", err)
		time.Sleep(1 * time.Second)
		continue
	}
}

func (p *EveFileProcessor) addCustomFields(event eve.EveEvent) {
	for key := range p.customFields {
		event[key] = p.customFields[key]
	}
}
<|endoftext|>


<commit_before>package statsd

import (
	"os"
	"strings"
	"time"
)

type options struct {
	timeout       time.Duration
	flushPeriod   time.Duration
	maxPacketSize int
	errHandler    func(error)

	prefix   string
	hostname string
}

type Option func(*options)

func ErrorHandler(h func(error)) Option {
	return func(o *options) {
		o.errHandler = h
	}
}

func FlushPeriod(d time.Duration) Option {
	return func(o *options) {
		o.flushPeriod = d
	}
}

func Prefix(s string) Option {
	return func(o *options) {
		o.prefix = s
	}
}

func MaxPacketSize(n int) Option {
	return func(o *options) {
		o.maxPacketSize = n
	}
}

func Timeout(d time.Duration) Option {
	return func(o *options) {
		o.timeout = d
	}
}

func Hostname(hostname string) Option {
	return func(o *options) {
		o.hostname = hostname
	}
}

type Client struct {
	opts options

	cc *clientConn
}

func New(network, addr string, opt ...Option) (*Client, error) {
	c := &Client{}
	for _, o := range opt {
		o(&c.opts)
	}

	if c.opts.hostname == "" {
		hostname, _ := os.Hostname()
		c.opts.hostname = strings.Replace(hostname, ".", "_", -1)
	}
	if c.opts.timeout <= 0 {
		c.opts.timeout = time.Second * 5
	}
	if c.opts.flushPeriod <= 0 {
		c.opts.flushPeriod = time.Millisecond * 100
	}
	if c.opts.maxPacketSize <= 0 {
		c.opts.maxPacketSize = 1400
	}

	cc, err := newClientConn(network, addr, c)
	if err != nil {
		return nil, err
	}

	c.cc = cc

	return c, nil
}

func (c *Client) Increment(bucket ...Field) {
	c.CountInt32(1, bucket...)
}

func (c *Client) CountInt32(n int32, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Int32(n), bucket))
}

func (c *Client) CountUint32(n uint32, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Uint32(n), bucket))
}

func (c *Client) CountInt64(n int64, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Int64(n), bucket))
}

func (c *Client) CountUint64(n uint64, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Uint64(n), bucket))
}

func (c *Client) GaugeInt32(n int32, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Int32(n), bucket))
}

func (c *Client) GaugeUint32(n uint32, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Uint32(n), bucket))
}

func (c *Client) GaugeInt64(n int64, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Int64(n), bucket))
}

func (c *Client) GaugeUint64(n uint64, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Uint64(n), bucket))
}

func (c *Client) TimingSince(start time.Time, bucket ...Field) {
	c.send(c.encode(MetricTypeTiming, Float64(float64(time.Now().Sub(start).Nanoseconds())/float64(time.Millisecond)), bucket))
}

func (c *Client) Timing(duration time.Duration, bucket ...Field) {
	c.send(c.encode(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), bucket))
}

func (c *Client) Incrementf(template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Int32(1), template, args))
}

func (c *Client) CountInt32f(n int32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Int32(n), template, args))
}

func (c *Client) CountUint32f(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Uint32(n), template, args))
}

func (c *Client) CountInt64f(n int64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Int64(n), template, args))
}

func (c *Client) CountUint64f(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Uint64(n), template, args))
}

func (c *Client) GaugeInt32f(n int32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Int32(n), template, args))
}

func (c *Client) GaugeUint32f(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Uint32(n), template, args))
}

func (c *Client) GaugeInt64f(n int64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Int64(n), template, args))
}

func (c *Client) GaugeUint64f(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Uint64(n), template, args))
}

func (c *Client) Timingf(duration time.Duration, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), template, args))
}

func (c *Client) TimingSincef(start time.Time, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeTiming, Float64(float64(time.Now().Sub(start).Nanoseconds())/float64(time.Millisecond)), template, args))
}

func (c *Client) IncrementWithHost(bucket ...Field) {
	c.CountInt32WithHost(1, bucket...)
}

func (c *Client) CountInt32WithHost(n int32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Int32(n), bucket))
}

func (c *Client) CountUint32WithHost(n uint32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Uint32(n), bucket))
}

func (c *Client) CountInt64WithHost(n int64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Int64(n), bucket))
}

func (c *Client) CountUint64WithHost(n uint64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Uint64(n), bucket))
}

func (c *Client) GaugeInt32WithHost(n int32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Int32(n), bucket))
}

func (c *Client) GaugeUint32WithHost(n uint32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Uint32(n), bucket))
}

func (c *Client) GaugeInt64WithHost(n int64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Int64(n), bucket))
}

func (c *Client) GaugeUint64WithHost(n uint64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Uint64(n), bucket))
}

func (c *Client) TimingSinceWithHost(start time.Time, bucket ...Field) {
	elapsed := float64(time.Now().Sub(start).Nanoseconds()) / float64(time.Millisecond)
	c.send(c.encodeWithHost(MetricTypeTiming, Float64(elapsed), bucket))
}

func (c *Client) TimingWithHost(duration time.Duration, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), bucket))
}

func (c *Client) IncrementfWithHost(template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Int32(1), template, args))
}

func (c *Client) CountInt32fWithHost(n int32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Int32(n), template, args))
}

func (c *Client) CountUint32fWithHost(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Uint32(n), template, args))
}

func (c *Client) CountInt64fWithHost(n int64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Int64(n), template, args))
}

func (c *Client) CountUint64fWithHost(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Uint64(n), template, args))
}

func (c *Client) GaugeInt32fWithHost(n int32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Int32(n), template, args))
}

func (c *Client) GaugeUint32fWithHost(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Uint32(n), template, args))
}

func (c *Client) GaugeInt64fWithHost(n int64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Int64(n), template, args))
}

func (c *Client) GaugeUint64fWithHost(n uint64, template string, args ...interface{}) {
	b := c.encodeTplWithHost(MetricTypeGauge, Uint64(n), template, args)
	c.send(b)
}

func (c *Client) TimingfWithHost(duration time.Duration, template string, args ...interface{}) {
	b := c.encodeTplWithHost(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), template, args)
	c.send(b)
}

func (c *Client) TimingSincefWithHost(start time.Time, template string, args ...interface{}) {
	elapsed := float64(time.Now().Sub(start).Nanoseconds()) / float64(time.Millisecond)
	b := c.encodeTplWithHost(MetricTypeTiming, Float64(elapsed), template, args)
	c.send(b)
}

func (c *Client) encode(typ MetricType, val Field, bucket []Field) *buf {
	return encode(typ, val, c.opts.prefix, "", bucket)
}

func (c *Client) encodeWithHost(typ MetricType, val Field, bucket []Field) *buf {
	return encode(typ, val, c.opts.prefix, c.opts.hostname, bucket)
}

func (c *Client) encodeTpl(typ MetricType, val Field, template string, fmtArgs []interface{}) *buf {
	return encodeTpl(typ, val, c.opts.prefix, "", template, fmtArgs)
}

func (c *Client) encodeTplWithHost(typ MetricType, val Field, template string, fmtArgs []interface{}) *buf {
	return encodeTpl(typ, val, c.opts.prefix, c.opts.hostname, template, fmtArgs)
}

func (c *Client) send(b *buf) {
	if b == nil {
		return
	}
	c.cc.write(b.Bytes())
	freeBuf(b)
}
<commit_msg>Support gauge float64<commit_after>package statsd

import (
	"os"
	"strings"
	"time"
)

type options struct {
	timeout       time.Duration
	flushPeriod   time.Duration
	maxPacketSize int
	errHandler    func(error)

	prefix   string
	hostname string
}

type Option func(*options)

func ErrorHandler(h func(error)) Option {
	return func(o *options) {
		o.errHandler = h
	}
}

func FlushPeriod(d time.Duration) Option {
	return func(o *options) {
		o.flushPeriod = d
	}
}

func Prefix(s string) Option {
	return func(o *options) {
		o.prefix = s
	}
}

func MaxPacketSize(n int) Option {
	return func(o *options) {
		o.maxPacketSize = n
	}
}

func Timeout(d time.Duration) Option {
	return func(o *options) {
		o.timeout = d
	}
}

func Hostname(hostname string) Option {
	return func(o *options) {
		o.hostname = hostname
	}
}

type Client struct {
	opts options

	cc *clientConn
}

func New(network, addr string, opt ...Option) (*Client, error) {
	c := &Client{}
	for _, o := range opt {
		o(&c.opts)
	}

	if c.opts.hostname == "" {
		hostname, _ := os.Hostname()
		c.opts.hostname = strings.Replace(hostname, ".", "_", -1)
	}
	if c.opts.timeout <= 0 {
		c.opts.timeout = time.Second * 5
	}
	if c.opts.flushPeriod <= 0 {
		c.opts.flushPeriod = time.Millisecond * 100
	}
	if c.opts.maxPacketSize <= 0 {
		c.opts.maxPacketSize = 1400
	}

	cc, err := newClientConn(network, addr, c)
	if err != nil {
		return nil, err
	}

	c.cc = cc

	return c, nil
}

func (c *Client) Increment(bucket ...Field) {
	c.CountInt32(1, bucket...)
}

func (c *Client) CountInt32(n int32, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Int32(n), bucket))
}

func (c *Client) CountUint32(n uint32, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Uint32(n), bucket))
}

func (c *Client) CountInt64(n int64, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Int64(n), bucket))
}

func (c *Client) CountUint64(n uint64, bucket ...Field) {
	c.send(c.encode(MetricTypeCount, Uint64(n), bucket))
}

func (c *Client) GaugeInt32(n int32, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Int32(n), bucket))
}

func (c *Client) GaugeUint32(n uint32, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Uint32(n), bucket))
}

func (c *Client) GaugeInt64(n int64, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Int64(n), bucket))
}

func (c *Client) GaugeUint64(n uint64, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Uint64(n), bucket))
}

func (c *Client) GaugeFloat64(n float64, bucket ...Field) {
	c.send(c.encode(MetricTypeGauge, Float64(n), bucket))
}

func (c *Client) TimingSince(start time.Time, bucket ...Field) {
	c.send(c.encode(MetricTypeTiming, Float64(float64(time.Now().Sub(start).Nanoseconds())/float64(time.Millisecond)), bucket))
}

func (c *Client) Timing(duration time.Duration, bucket ...Field) {
	c.send(c.encode(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), bucket))
}

func (c *Client) Incrementf(template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Int32(1), template, args))
}

func (c *Client) CountInt32f(n int32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Int32(n), template, args))
}

func (c *Client) CountUint32f(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Uint32(n), template, args))
}

func (c *Client) CountInt64f(n int64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Int64(n), template, args))
}

func (c *Client) CountUint64f(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeCount, Uint64(n), template, args))
}

func (c *Client) GaugeInt32f(n int32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Int32(n), template, args))
}

func (c *Client) GaugeUint32f(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Uint32(n), template, args))
}

func (c *Client) GaugeInt64f(n int64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Int64(n), template, args))
}

func (c *Client) GaugeUint64f(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Uint64(n), template, args))
}

func (c *Client) GaugeFloat64f(n float64, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeGauge, Float64(n), template, args))
}

func (c *Client) Timingf(duration time.Duration, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), template, args))
}

func (c *Client) TimingSincef(start time.Time, template string, args ...interface{}) {
	c.send(c.encodeTpl(MetricTypeTiming, Float64(float64(time.Now().Sub(start).Nanoseconds())/float64(time.Millisecond)), template, args))
}

func (c *Client) IncrementWithHost(bucket ...Field) {
	c.CountInt32WithHost(1, bucket...)
}

func (c *Client) CountInt32WithHost(n int32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Int32(n), bucket))
}

func (c *Client) CountUint32WithHost(n uint32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Uint32(n), bucket))
}

func (c *Client) CountInt64WithHost(n int64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Int64(n), bucket))
}

func (c *Client) CountUint64WithHost(n uint64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeCount, Uint64(n), bucket))
}

func (c *Client) GaugeInt32WithHost(n int32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Int32(n), bucket))
}

func (c *Client) GaugeUint32WithHost(n uint32, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Uint32(n), bucket))
}

func (c *Client) GaugeInt64WithHost(n int64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Int64(n), bucket))
}

func (c *Client) GaugeUint64WithHost(n uint64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Uint64(n), bucket))
}

func (c *Client) GaugeFloat64WithHost(n float64, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeGauge, Float64(n), bucket))
}

func (c *Client) TimingSinceWithHost(start time.Time, bucket ...Field) {
	elapsed := float64(time.Now().Sub(start).Nanoseconds()) / float64(time.Millisecond)
	c.send(c.encodeWithHost(MetricTypeTiming, Float64(elapsed), bucket))
}

func (c *Client) TimingWithHost(duration time.Duration, bucket ...Field) {
	c.send(c.encodeWithHost(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), bucket))
}

func (c *Client) IncrementfWithHost(template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Int32(1), template, args))
}

func (c *Client) CountInt32fWithHost(n int32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Int32(n), template, args))
}

func (c *Client) CountUint32fWithHost(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Uint32(n), template, args))
}

func (c *Client) CountInt64fWithHost(n int64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Int64(n), template, args))
}

func (c *Client) CountUint64fWithHost(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeCount, Uint64(n), template, args))
}

func (c *Client) GaugeInt32fWithHost(n int32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Int32(n), template, args))
}

func (c *Client) GaugeUint32fWithHost(n uint32, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Uint32(n), template, args))
}

func (c *Client) GaugeInt64fWithHost(n int64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Int64(n), template, args))
}

func (c *Client) GaugeUint64fWithHost(n uint64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Uint64(n), template, args))
}

func (c *Client) GaugeFloat64fWithHost(n float64, template string, args ...interface{}) {
	c.send(c.encodeTplWithHost(MetricTypeGauge, Float64(n), template, args))
}

func (c *Client) TimingfWithHost(duration time.Duration, template string, args ...interface{}) {
	b := c.encodeTplWithHost(MetricTypeTiming, Float64(float64(duration)/float64(time.Millisecond)), template, args)
	c.send(b)
}

func (c *Client) TimingSincefWithHost(start time.Time, template string, args ...interface{}) {
	elapsed := float64(time.Now().Sub(start).Nanoseconds()) / float64(time.Millisecond)
	b := c.encodeTplWithHost(MetricTypeTiming, Float64(elapsed), template, args)
	c.send(b)
}

func (c *Client) encode(typ MetricType, val Field, bucket []Field) *buf {
	return encode(typ, val, c.opts.prefix, "", bucket)
}

func (c *Client) encodeWithHost(typ MetricType, val Field, bucket []Field) *buf {
	return encode(typ, val, c.opts.prefix, c.opts.hostname, bucket)
}

func (c *Client) encodeTpl(typ MetricType, val Field, template string, fmtArgs []interface{}) *buf {
	return encodeTpl(typ, val, c.opts.prefix, "", template, fmtArgs)
}

func (c *Client) encodeTplWithHost(typ MetricType, val Field, template string, fmtArgs []interface{}) *buf {
	return encodeTpl(typ, val, c.opts.prefix, c.opts.hostname, template, fmtArgs)
}

func (c *Client) send(b *buf) {
	if b == nil {
		return
	}
	c.cc.write(b.Bytes())
	freeBuf(b)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"reflect"
	"strings"
)

const (
	networkDeviceStatsMetricPrefix = "net.dev."
	netstatStatsMetricPrefix       = "net.netstat."
)

func collectNetworkDeviceStats() {
	stats, err := readNetworkDeviceStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't collect network device stats.")
	}

	for _, stat := range stats {
		if ifaceRegExp == nil || !ifaceRegExp.MatchString(stat.Iface) {
			elem := reflect.ValueOf(&stat).Elem()
			typeOfElem := elem.Type()

			for i := 0; i < elem.NumField(); i++ {
				field := typeOfElem.Field(i)

				if field.Name == "Iface" {
					continue
				}

				value := elem.Field(i).Uint()

				_, collect := StatsMap[field.Name]

				if metricName := field.Tag.Get("json"); collect && metricName != "" {
					if err := StatsdClient.Count(strings.Join([]string{networkDeviceStatsMetricPrefix, metricName}, ""), int64(value), []string{"iface:" + stat.Iface}, 1); err != nil {
						Log.WithField("error", err).Error("Couldn't submit event to statsd.")
					}
				}
			}
		} else {
			Log.Infof("Filtered interface %s from network device stat collection.", stat.Iface)
		}
	}
}

func logNetworkDeviceStats() {
	stats, err := readNetworkDeviceStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't log network device stats.")
	}

	jsonBytes, err := json.Marshal(stats)
	if err != nil {
		Log.WithField("error", err).Errorf("Error JSON marshalling: %+v.", stats)
	}

	Log.WithField("stat", NetworkStatPath).Infof(string(jsonBytes))
}

func collectNetstatStats() {
	stats, err := readNetstatStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't collect netstat stats.")
	}

	elem := reflect.ValueOf(stats).Elem()
	typeOfElem := elem.Type()

	for i := 0; i < elem.NumField(); i++ {
		field := typeOfElem.Field(i)

		value := elem.Field(i).Uint()

		_, collect := StatsMap[field.Name]

		if metricName := field.Tag.Get("json"); collect && metricName != "" {
			if err := StatsdClient.Count(strings.Join([]string{netstatStatsMetricPrefix, metricName}, ""), int64(value), []string{}, 1); err != nil {
				Log.WithField("error", err).Error("Couldn't submit event to statsd.")
			}
		}
	}
}

func logNetstatStats() {
	stats, err := readNetstatStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't log netstat stats.")
	}

	jsonBytes, err := json.Marshal(stats)
	if err != nil {
		Log.WithField("error", err).Errorf("Error JSON marshalling: %+v.", stats)
	}

	Log.WithField("stat", NetstatStatPath).Infof(string(jsonBytes))
}
<commit_msg>Cleaner stat logging (without escape characters).<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"time"
)

const (
	networkDeviceStatsMetricPrefix = "net.dev."
	netstatStatsMetricPrefix       = "net.netstat."
)

type LoggedStat struct {
	Name string      `json:"name"`
	Time time.Time   `json:"time"`
	Stat interface{} `json:"stat"`
}

func collectNetworkDeviceStats() {
	stats, err := readNetworkDeviceStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't collect network device stats.")
	}

	for _, stat := range stats {
		if ifaceRegExp == nil || !ifaceRegExp.MatchString(stat.Iface) {
			elem := reflect.ValueOf(&stat).Elem()
			typeOfElem := elem.Type()

			for i := 0; i < elem.NumField(); i++ {
				field := typeOfElem.Field(i)

				if field.Name == "Iface" {
					continue
				}

				value := elem.Field(i).Uint()

				_, collect := StatsMap[field.Name]

				if metricName := field.Tag.Get("json"); collect && metricName != "" {
					if err := StatsdClient.Count(strings.Join([]string{networkDeviceStatsMetricPrefix, metricName}, ""), int64(value), []string{"iface:" + stat.Iface}, 1); err != nil {
						Log.WithField("error", err).Error("Couldn't submit event to statsd.")
					}
				}
			}
		} else {
			Log.Infof("Filtered interface %s from network device stat collection.", stat.Iface)
		}
	}
}

func logNetworkDeviceStats() {
	stats, err := readNetworkDeviceStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't log network device stats.")
	}

	loggedStat := LoggedStat{NetworkStatPath, time.Now(), stats}

	jsonBytes, err := json.Marshal(loggedStat)
	if err != nil {
		Log.WithField("error", err).Errorf("Error JSON marshalling: %+v.", stats)
	}

	fmt.Println(string(jsonBytes))
}

func collectNetstatStats() {
	stats, err := readNetstatStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't collect netstat stats.")
	}

	elem := reflect.ValueOf(stats).Elem()
	typeOfElem := elem.Type()

	for i := 0; i < elem.NumField(); i++ {
		field := typeOfElem.Field(i)

		value := elem.Field(i).Uint()

		_, collect := StatsMap[field.Name]

		if metricName := field.Tag.Get("json"); collect && metricName != "" {
			if err := StatsdClient.Count(strings.Join([]string{netstatStatsMetricPrefix, metricName}, ""), int64(value), []string{}, 1); err != nil {
				Log.WithField("error", err).Error("Couldn't submit event to statsd.")
			}
		}
	}
}

func logNetstatStats() {
	stats, err := readNetstatStats()

	if err != nil {
		Log.WithField("error", err).Error("Error reading network device stats. Can't log netstat stats.")
	}

	loggedStat := LoggedStat{NetstatStatPath, time.Now(), stats}

	jsonBytes, err := json.Marshal(loggedStat)
	if err != nil {
		Log.WithField("error", err).Errorf("Error JSON marshalling: %+v.", stats)
	}

	fmt.Println(string(jsonBytes))
}
<|endoftext|>


<commit_before>// statspipe is a metrics pipeline
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	//"regexp"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	//"github.com/davecgh/go-spew/spew"
	"github.com/davecheney/profile"
)

const FlushInterval = time.Duration(10 * time.Second)
const BufSize = 8192

//-----------------------------------------------------------------------------

// Command line flags
var (
	listen   = flag.String("listen", ":1514", "Listener address")
	graphite = flag.String("graphite", "localhost:2003", "Graphite server address")

	// Profiling
	cpuprofile   = flag.Bool("cpuprofile", false, "Enable CPU profiling")
	memprofile   = flag.Bool("memprofile", false, "Enable memory profiling")
	blockprofile = flag.Bool("blockprofile", false, "Enable block profiling")
)

//-----------------------------------------------------------------------------
// Data structures

// Metric is a numeric data point
type Metric struct {
	Bucket string
	Value  interface{}
	Type   string
}

// Metrics should be in statsd format. Metric names may not have spaces.
//
//     <metric_name>:<metric_value>|<metric_type>|@<sample_rate>
//
// Note: The sample rate is optional
// var statsPattern = regexp.MustCompile(`[\w\.]+:-?\d+\|(?:c|ms|g)(?:\|\@[\d\.]+)?`)

// In is a channel for processing metrics
var In = make(chan *Metric)

// counters holds all of the counter metrics
var counters = struct {
	sync.RWMutex
	m map[string]int64
}{m: make(map[string]int64)}

// gauges holds all of the gauge metrics
var gauges = struct {
	sync.RWMutex
	m map[string]uint64
}{m: make(map[string]uint64)}

// Timers is a list of uints
type Timers []uint64

// timers holds all of the timer metrics
var timers = struct {
	sync.RWMutex
	m map[string]Timers
}{m: make(map[string]Timers)}

// Internal metrics
type Stats struct {
	IngressRate     int64
	IngressMetrics  int64
	IngressCounters int64
	IngressGauges   int64
	IngressTimers   int64
}

var stats = &Stats{}

// TODO: move this to command line option
var Percentiles = []int{5, 95}

//-----------------------------------------------------------------------------

// Implement the sort interface for Timers
func (t Timers) Len() int           { return len(t) }
func (t Timers) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
func (t Timers) Less(i, j int) bool { return t[i] < t[j] }

//-----------------------------------------------------------------------------

// ListenUDP creates a UDP listener
func ListenUDP(addr string) error {
	var buf = make([]byte, 1024)
	ln, err := net.ResolveUDPAddr("udp", addr)

	if err != nil {
		return err
	}

	sock, err := net.ListenUDP("udp", ln)

	if err != nil {
		return err
	}

	log.Printf("Listening on UDP %s\n", ln)

	for {
		n, raddr, err := sock.ReadFromUDP(buf[:])

		if err != nil {
			// TODO: handle error
			continue
		}

		log.Printf("Read %d bytes from %s\n", n, raddr)
		go handleMessage(buf)
	}
}

// ListenTCP creates a TCP listener
func ListenTCP(addr string) error {
	l, err := net.Listen("tcp", addr)

	if err != nil {
		return err
	}

	defer l.Close()
	log.Printf("Listening on TCP %s\n", l.Addr())

	for {
		conn, err := l.Accept()

		if err != nil {
			// TODO: handle error
			continue
		}

		go handleConnection(conn)
	}
}

// handleConnection handles a single client connection
func handleConnection(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)

	// Incoming metrics should be separated by a newline
	for {
		line, err := r.ReadBytes('\n')

		if err != nil {
			if err == io.EOF {
				break
			} else {
				// TODO: handle error
			}
		}

		handleMessage(line)
	}
}

// Handle an event message
func handleMessage(buf []byte) {
	// According to the statsd protocol, metrics should be separated by a
	// newline. This parser isn't quite as strict since it may be receiving
	// metrics from clients that aren't proper statsd clients (e.g. syslog).
	// Instead, find metrics by looking for tokens between spaces and then
	// use parseMetric to create an actual metric that can be aggregated.
	tokens := bytes.Split(buf, []byte(" "))

	for _, token := range tokens {
		// metrics must have a : and | at a minimum
		if bytes.Index(token, []byte(":")) == -1 ||
			bytes.Index(token, []byte("|")) == -1 {
			continue
		}

		metric, err := parseMetric(token)

		if err != nil {
			log.Printf("ERROR: Unable to parse metric %s: %s",
				token, err)
			continue
		}

		// Send metric off for processing
		In <- metric
	}
}

// parseMetric parses a raw metric into a Metric struct
func parseMetric(b []byte) (*Metric, error) {
	i := bytes.Index(b, []byte(":"))
	j := bytes.Index(b, []byte("|"))
	k := bytes.Index(b, []byte("@"))
	v := b[i+1 : j]

	// End position of the metric type is the end of the byte slice
	// if no sample was sent.
	tEnd := len(b)
	var sampleRate float64 = 1

	// Indicates that a sample rate was sent as part of the metric
	if k > -1 {
		tEnd = k - 1 // Use -1 because of the | before the @
		sr := b[(k + 1):len(b)]
		var err error
		sampleRate, err = strconv.ParseFloat(string(sr), 64)

		if err != nil {
			return nil, err
		}
	}

	m := &Metric{
		Bucket: string(b[0:i]),
		Type:   string(b[j+1 : tEnd]),
	}

	switch m.Type {
	case "c":
		val, err := strconv.ParseInt(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = int64(float64(val) / sampleRate)
	default:
		val, err := strconv.ParseUint(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = val
	}

	return m, nil
}

// processMetrics updates new metrics and flushes aggregates to Graphite
func processMetrics() {
	ticker := time.NewTicker(FlushInterval)

	for {
		select {
		case <-ticker.C:
			flushMetrics()
		case m := <-In:
			atomic.AddInt64(&stats.IngressMetrics, 1)

			switch m.Type {
			case "c":
				counters.Lock()
				counters.m[m.Bucket] += m.Value.(int64)
				counters.Unlock()
				atomic.AddInt64(&stats.IngressCounters, 1)

			case "g":
				gauges.Lock()
				gauges.m[m.Bucket] = m.Value.(uint64)
				gauges.Unlock()
				atomic.AddInt64(&stats.IngressGauges, 1)

			case "ms":
				timers.Lock()
				_, ok := timers.m[m.Bucket]

				if !ok {
					var t Timers
					timers.m[m.Bucket] = t
				}

				timers.m[m.Bucket] = append(timers.m[m.Bucket], m.Value.(uint64))
				timers.Unlock()
				atomic.AddInt64(&stats.IngressTimers, 1)

			}
		}
	}
}

// flushMetrics sends metrics to Graphite
func flushMetrics() {
	var buf bytes.Buffer
	now := time.Now().Unix()

	log.Printf("%+v", stats)

	// Build buffer of stats
	flushCounters(&buf, now)
	flushGauges(&buf, now)
	flushTimers(&buf, now)
	flushInternalStats(&buf, now)

	// Send metrics to Graphite
	sendGraphite(&buf)
}

// flushInternalStats writes the internal stats to the buffer
func flushInternalStats(buf *bytes.Buffer, now int64) {
	//fmt.Fprintf(buf, "statsd.metrics.per_second %d %d\n", v, now)
	fmt.Fprintf(buf, "statsd.metrics.count %d %d\n",
		atomic.LoadInt64(&stats.IngressMetrics), now)
	fmt.Fprintf(buf, "statsd.counters.count %d %d\n",
		atomic.LoadInt64(&stats.IngressCounters), now)
	fmt.Fprintf(buf, "statsd.gauges.count %d %d\n",
		atomic.LoadInt64(&stats.IngressGauges), now)
	fmt.Fprintf(buf, "statsd.timers.count %d %d\n",
		atomic.LoadInt64(&stats.IngressTimers), now)

	// Clear internal metrics
	atomic.StoreInt64(&stats.IngressMetrics, 0)
	atomic.StoreInt64(&stats.IngressCounters, 0)
	atomic.StoreInt64(&stats.IngressGauges, 0)
	atomic.StoreInt64(&stats.IngressTimers, 0)
}

// flushCounters writes the counters to the buffer
func flushCounters(buf *bytes.Buffer, now int64) {
	counters.Lock()
	defer counters.Unlock()

	for k, v := range counters.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(counters.m, k)
	}
}

// flushGauges writes the gauges to the buffer
func flushGauges(buf *bytes.Buffer, now int64) {
	gauges.Lock()
	defer gauges.Unlock()

	for k, v := range gauges.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(gauges.m, k)
	}
}

// flushTimers writes the timers and aggregate statistics to the buffer
func flushTimers(buf *bytes.Buffer, now int64) {
	timers.RLock()
	defer timers.RUnlock()
	var n int64

	for k, t := range timers.m {
		count := len(t)

		// Skip processing if there are no timer values
		if count < 1 {
			break
		}

		var sum uint64

		for _, v := range t {
			sum += v
			n++
		}

		// Linear average (mean)
		mean := float64(sum) / float64(count)

		// Min and Max
		sort.Sort(t)
		min := t[0]
		max := t[len(t)-1]

		// Write out all derived stats
		fmt.Fprintf(buf, "%s.count %d %d\n", k, count, now)
		fmt.Fprintf(buf, "%s.mean %f %d\n", k, mean, now)
		fmt.Fprintf(buf, "%s.lower %d %d\n", k, min, now)
		fmt.Fprintf(buf, "%s.upper %d %d\n", k, max, now)

		// Calculate and write out percentiles
		for _, pct := range Percentiles {
			p := perc(t, pct)
			fmt.Fprintf(buf, "%s.perc%d %f %d\n", k, pct, p, now)
		}

		delete(timers.m, k)
	}
}

// percentile calculates Nth percentile of a list of values
func perc(values []uint64, pct int) float64 {
	p := float64(pct) / float64(100)
	n := float64(len(values))
	i := math.Ceil(p*n) - 1

	return float64(values[int(i)])
}

// sendGraphite sends metrics to graphite
func sendGraphite(buf *bytes.Buffer) {
	conn, err := net.Dial("tcp", *graphite)

	if err != nil {
		log.Fatal("ERROR: Unable to connect to graphite")
	}

	w := bufio.NewWriter(conn)
	n, err := buf.WriteTo(w)

	if err != nil {
		log.Fatal("ERROR: Unable to write to graphite")
	}

	w.Flush()
	conn.Close()

	log.Printf("Wrote %d bytes to Graphite", n)
}

//-----------------------------------------------------------------------------

func main() {
	flag.Parse()

	// Profiling
	cfg := profile.Config{
		CPUProfile:   *cpuprofile,
		MemProfile:   *memprofile,
		BlockProfile: *blockprofile,
		ProfilePath:  ".",
	}

	p := profile.Start(&cfg)
	defer p.Stop()

	// Process metrics as they arrive
	go processMetrics()

	// Setup listeners
	go log.Fatal(ListenTCP(*listen))
}
<commit_msg>Fix UDP listener to split on newline before sending to handleMessage.<commit_after>// statspipe is a metrics pipeline
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	//"regexp"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	//"github.com/davecgh/go-spew/spew"
	"github.com/davecheney/profile"
)

const FlushInterval = time.Duration(10 * time.Second)
const BufSize = 8192

//-----------------------------------------------------------------------------

// Command line flags
var (
	listen   = flag.String("listen", ":8125", "Listener address")
	graphite = flag.String("graphite", "localhost:2003", "Graphite server address")

	// Profiling
	cpuprofile   = flag.Bool("cpuprofile", false, "Enable CPU profiling")
	memprofile   = flag.Bool("memprofile", false, "Enable memory profiling")
	blockprofile = flag.Bool("blockprofile", false, "Enable block profiling")
)

//-----------------------------------------------------------------------------
// Data structures

// Metric is a numeric data point
type Metric struct {
	Bucket string
	Value  interface{}
	Type   string
}

// Metrics should be in statsd format. Metric names may not have spaces.
//
//     <metric_name>:<metric_value>|<metric_type>|@<sample_rate>
//
// Note: The sample rate is optional
// var statsPattern = regexp.MustCompile(`[\w\.]+:-?\d+\|(?:c|ms|g)(?:\|\@[\d\.]+)?`)

// In is a channel for processing metrics
var In = make(chan *Metric)

// counters holds all of the counter metrics
var counters = struct {
	sync.RWMutex
	m map[string]int64
}{m: make(map[string]int64)}

// gauges holds all of the gauge metrics
var gauges = struct {
	sync.RWMutex
	m map[string]uint64
}{m: make(map[string]uint64)}

// Timers is a list of uints
type Timers []uint64

// timers holds all of the timer metrics
var timers = struct {
	sync.RWMutex
	m map[string]Timers
}{m: make(map[string]Timers)}

// Internal metrics
type Stats struct {
	IngressRate     int64
	IngressMetrics  int64
	IngressCounters int64
	IngressGauges   int64
	IngressTimers   int64
}

var stats = &Stats{}

// TODO: move this to command line option
var Percentiles = []int{5, 95}

//-----------------------------------------------------------------------------

// Implement the sort interface for Timers
func (t Timers) Len() int           { return len(t) }
func (t Timers) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
func (t Timers) Less(i, j int) bool { return t[i] < t[j] }

//-----------------------------------------------------------------------------

// ListenUDP creates a UDP listener
func ListenUDP(addr string) error {
	var buf = make([]byte, 1024)
	ln, err := net.ResolveUDPAddr("udp", addr)

	if err != nil {
		return err
	}

	sock, err := net.ListenUDP("udp", ln)

	if err != nil {
		return err
	}

	log.Printf("Listening on UDP %s\n", ln)

	for {
		n, raddr, err := sock.ReadFromUDP(buf[:])

		if err != nil {
			// TODO: handle error
			continue
		}

		go handleUdpMessage(buf)
	}
}

func handleUdpMessage(buf []byte) {
	tokens := bytes.Split(buf, []byte("\n"))

	for _, token := range tokens {
		handleMessage(token)
	}
}

// ListenTCP creates a TCP listener
func ListenTCP(addr string) error {
	l, err := net.Listen("tcp", addr)

	if err != nil {
		return err
	}

	defer l.Close()
	log.Printf("Listening on TCP %s\n", l.Addr())

	for {
		conn, err := l.Accept()

		if err != nil {
			// TODO: handle error
			continue
		}

		go handleConnection(conn)
	}
}

// handleConnection handles a single client connection
func handleConnection(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)

	// Incoming metrics should be separated by a newline
	for {
		line, err := r.ReadBytes('\n')

		if err != nil {
			if err == io.EOF {
				break
			} else {
				// TODO: handle error
			}
		}

		handleMessage(line)
	}
}

// Handle an event message
func handleMessage(buf []byte) {
	// According to the statsd protocol, metrics should be separated by a
	// newline. This parser isn't quite as strict since it may be receiving
	// metrics from clients that aren't proper statsd clients (e.g. syslog).
	// Instead, find metrics by looking for tokens between spaces and then
	// use parseMetric to create an actual metric that can be aggregated.
	tokens := bytes.Split(buf, []byte(" "))

	for _, token := range tokens {
		// metrics must have a : and | at a minimum
		if bytes.Index(token, []byte(":")) == -1 ||
			bytes.Index(token, []byte("|")) == -1 {
			continue
		}

		metric, err := parseMetric(token)

		if err != nil {
			log.Printf("ERROR: Unable to parse metric %s: %s",
				token, err)
			continue
		}

		// Send metric off for processing
		In <- metric
	}
}

// parseMetric parses a raw metric into a Metric struct
func parseMetric(b []byte) (*Metric, error) {
	i := bytes.Index(b, []byte(":"))
	j := bytes.Index(b, []byte("|"))
	k := bytes.Index(b, []byte("@"))
	v := b[i+1 : j]

	// End position of the metric type is the end of the byte slice
	// if no sample was sent.
	tEnd := len(b)
	var sampleRate float64 = 1

	// Indicates that a sample rate was sent as part of the metric
	if k > -1 {
		tEnd = k - 1 // Use -1 because of the | before the @
		sr := b[(k + 1):len(b)]
		var err error
		sampleRate, err = strconv.ParseFloat(string(sr), 64)

		if err != nil {
			return nil, err
		}
	}

	m := &Metric{
		Bucket: string(b[0:i]),
		Type:   string(b[j+1 : tEnd]),
	}

	switch m.Type {
	case "c":
		val, err := strconv.ParseInt(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = int64(float64(val) / sampleRate)
	default:
		val, err := strconv.ParseUint(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = val
	}

	return m, nil
}

// processMetrics updates new metrics and flushes aggregates to Graphite
func processMetrics() {
	ticker := time.NewTicker(FlushInterval)

	for {
		select {
		case <-ticker.C:
			flushMetrics()
		case m := <-In:
			atomic.AddInt64(&stats.IngressMetrics, 1)

			switch m.Type {
			case "c":
				counters.Lock()
				counters.m[m.Bucket] += m.Value.(int64)
				counters.Unlock()
				atomic.AddInt64(&stats.IngressCounters, 1)

			case "g":
				gauges.Lock()
				gauges.m[m.Bucket] = m.Value.(uint64)
				gauges.Unlock()
				atomic.AddInt64(&stats.IngressGauges, 1)

			case "ms":
				timers.Lock()
				_, ok := timers.m[m.Bucket]

				if !ok {
					var t Timers
					timers.m[m.Bucket] = t
				}

				timers.m[m.Bucket] = append(timers.m[m.Bucket], m.Value.(uint64))
				timers.Unlock()
				atomic.AddInt64(&stats.IngressTimers, 1)

			}
		}
	}
}

// flushMetrics sends metrics to Graphite
func flushMetrics() {
	var buf bytes.Buffer
	now := time.Now().Unix()

	log.Printf("%+v", stats)

	// Build buffer of stats
	flushCounters(&buf, now)
	flushGauges(&buf, now)
	flushTimers(&buf, now)
	flushInternalStats(&buf, now)

	// Send metrics to Graphite
	sendGraphite(&buf)
}

// flushInternalStats writes the internal stats to the buffer
func flushInternalStats(buf *bytes.Buffer, now int64) {
	//fmt.Fprintf(buf, "statsd.metrics.per_second %d %d\n", v, now)
	fmt.Fprintf(buf, "statsd.metrics.count %d %d\n",
		atomic.LoadInt64(&stats.IngressMetrics), now)
	fmt.Fprintf(buf, "statsd.counters.count %d %d\n",
		atomic.LoadInt64(&stats.IngressCounters), now)
	fmt.Fprintf(buf, "statsd.gauges.count %d %d\n",
		atomic.LoadInt64(&stats.IngressGauges), now)
	fmt.Fprintf(buf, "statsd.timers.count %d %d\n",
		atomic.LoadInt64(&stats.IngressTimers), now)

	// Clear internal metrics
	atomic.StoreInt64(&stats.IngressMetrics, 0)
	atomic.StoreInt64(&stats.IngressCounters, 0)
	atomic.StoreInt64(&stats.IngressGauges, 0)
	atomic.StoreInt64(&stats.IngressTimers, 0)
}

// flushCounters writes the counters to the buffer
func flushCounters(buf *bytes.Buffer, now int64) {
	counters.Lock()
	defer counters.Unlock()

	for k, v := range counters.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(counters.m, k)
	}
}

// flushGauges writes the gauges to the buffer
func flushGauges(buf *bytes.Buffer, now int64) {
	gauges.Lock()
	defer gauges.Unlock()

	for k, v := range gauges.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(gauges.m, k)
	}
}

// flushTimers writes the timers and aggregate statistics to the buffer
func flushTimers(buf *bytes.Buffer, now int64) {
	timers.RLock()
	defer timers.RUnlock()
	var n int64

	for k, t := range timers.m {
		count := len(t)

		// Skip processing if there are no timer values
		if count < 1 {
			break
		}

		var sum uint64

		for _, v := range t {
			sum += v
			n++
		}

		// Linear average (mean)
		mean := float64(sum) / float64(count)

		// Min and Max
		sort.Sort(t)
		min := t[0]
		max := t[len(t)-1]

		// Write out all derived stats
		fmt.Fprintf(buf, "%s.count %d %d\n", k, count, now)
		fmt.Fprintf(buf, "%s.mean %f %d\n", k, mean, now)
		fmt.Fprintf(buf, "%s.lower %d %d\n", k, min, now)
		fmt.Fprintf(buf, "%s.upper %d %d\n", k, max, now)

		// Calculate and write out percentiles
		for _, pct := range Percentiles {
			p := perc(t, pct)
			fmt.Fprintf(buf, "%s.perc%d %f %d\n", k, pct, p, now)
		}

		delete(timers.m, k)
	}
}

// percentile calculates Nth percentile of a list of values
func perc(values []uint64, pct int) float64 {
	p := float64(pct) / float64(100)
	n := float64(len(values))
	i := math.Ceil(p*n) - 1

	return float64(values[int(i)])
}

// sendGraphite sends metrics to graphite
func sendGraphite(buf *bytes.Buffer) {
	conn, err := net.Dial("tcp", *graphite)

	if err != nil {
		log.Println("ERROR: Unable to connect to graphite")
		return
	}

	w := bufio.NewWriter(conn)
	n, err := buf.WriteTo(w)

	if err != nil {
		log.Println("ERROR: Unable to write to graphite")
	}

	w.Flush()
	conn.Close()

	log.Printf("Wrote %d bytes to Graphite", n)
}

//-----------------------------------------------------------------------------

func main() {
	flag.Parse()

	// Profiling
	cfg := profile.Config{
		CPUProfile:   *cpuprofile,
		MemProfile:   *memprofile,
		BlockProfile: *blockprofile,
		ProfilePath:  ".",
	}

	p := profile.Start(&cfg)
	defer p.Stop()

	// Process metrics as they arrive
	go processMetrics()

	// Setup listeners
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		log.Fatal(ListenUDP(*listen))
	}()

	go func() {
		defer wg.Done()
		log.Fatal(ListenTCP(*listen))
	}()

	wg.Wait()
}
<|endoftext|>


<commit_before>package ftp

// FTP status codes, defined in RFC 959
const (
	StatusInitiating    = 100
	StatusRestartMarker = 110
	StatusReadyMinute   = 120
	StatusAlreadyOpen   = 125
	StatusAboutToSend   = 150

	StatusCommandOK             = 200
	StatusCommandNotImplemented = 202
	StatusSystem                = 211
	StatusDirectory             = 212
	StatusFile                  = 213
	StatusHelp                  = 214
	StatusName                  = 215
	StatusReady                 = 220
	StatusClosing               = 221
	StatusDataConnectionOpen    = 225
	StatusClosingDataConnection = 226
	StatusPassiveMode           = 227
	StatusLongPassiveMode       = 228
	StatusExtendedPassiveMode   = 229
	StatusLoggedIn              = 230
	StatusLoggedOut             = 231
	StatusLogoutAck             = 232
	StatusRequestedFileActionOK = 250
	StatusPathCreated           = 257

	StatusUserOK             = 331
	StatusLoginNeedAccount   = 332
	StatusRequestFilePending = 350

	StatusNotAvailable             = 421
	StatusCanNotOpenDataConnection = 425
	StatusTransfertAborted         = 426
	StatusInvalidCredentials       = 430
	StatusHostUnavailable          = 434
	StatusFileActionIgnored        = 450
	StatusActionAborted            = 451
	Status452                      = 452

	StatusBadCommand              = 500
	StatusBadArguments            = 501
	StatusNotImplemented          = 502
	StatusBadSequence             = 503
	StatusNotImplementedParameter = 504
	StatusNotLoggedIn             = 530
	StatusStorNeedAccount         = 532
	StatusFileUnavailable         = 550
	StatusPageTypeUnknown         = 551
	StatusExceededStorage         = 552
	StatusBadFileName             = 553
)

var statusText = map[int]string{
	// 200
	StatusCommandOK:             "Command okay.",
	StatusCommandNotImplemented: "Command not implemented, superfluous at this site.",
	StatusSystem:                "System status, or system help reply.",
	StatusDirectory:             "Directory status.",
	StatusFile:                  "File status.",
	StatusHelp:                  "Help message.",
	StatusName:                  "",
	StatusReady:                 "Service ready for new user.",
	StatusClosing:               "Service closing control connection.",
	StatusDataConnectionOpen:    "Data connection open; no transfer in progress.",
	StatusClosingDataConnection: "Closing data connection. Requested file action successful.",
	StatusPassiveMode:           "Entering Passive Mode.",
	StatusLongPassiveMode:       "Entering Long Passive Mode.",
	StatusExtendedPassiveMode:   "Entering Extended Passive Mode.",
	StatusLoggedIn:              "User logged in, proceed.",
	StatusLoggedOut:             "User logged out; service terminated.",
	StatusLogoutAck:             "Logout command noted, will complete when transfer done.",
	StatusRequestedFileActionOK: "Requested file action okay, completed.",
	StatusPathCreated:           "Path created.",

	// 300
	StatusUserOK:             "User name okay, need password.",
	StatusLoginNeedAccount:   "Need account for login.",
	StatusRequestFilePending: "Requested file action pending further information.",

	// 400
	StatusNotAvailable:             "Service not available, closing control connection.",
	StatusCanNotOpenDataConnection: "Can't open data connection.",
	StatusTransfertAborted:         "Connection closed; transfer aborted.",
	StatusInvalidCredentials:       "Invalid username or password.",
	StatusHostUnavailable:          "Requested host unavailable.",
	StatusFileActionIgnored:        "Requested file action not taken.",
	StatusActionAborted:            "Requested action aborted. Local error in processing.",
	Status452:                      "Insufficient storage space in system.",

	// 500
	StatusBadCommand:              "Command unrecognized.",
	StatusBadArguments:            "Syntax error in parameters or arguments.",
	StatusNotImplemented:          "Command not implemented.",
	StatusBadSequence:             "Bad sequence of commands.",
	StatusNotImplementedParameter: "Command not implemented for that parameter.",
	StatusNotLoggedIn:             "Not logged in.",
	StatusStorNeedAccount:         "Need account for storing files.",
	StatusFileUnavailable:         "File unavailable.",
	StatusPageTypeUnknown:         "Page type unknown.",
	StatusExceededStorage:         "Exceeded storage allocation.",
	StatusBadFileName:             "File name not allowed.",
}
<commit_msg>StatusText returns a text for the FTP status code<commit_after>package ftp

// FTP status codes, defined in RFC 959
const (
	StatusInitiating    = 100
	StatusRestartMarker = 110
	StatusReadyMinute   = 120
	StatusAlreadyOpen   = 125
	StatusAboutToSend   = 150

	StatusCommandOK             = 200
	StatusCommandNotImplemented = 202
	StatusSystem                = 211
	StatusDirectory             = 212
	StatusFile                  = 213
	StatusHelp                  = 214
	StatusName                  = 215
	StatusReady                 = 220
	StatusClosing               = 221
	StatusDataConnectionOpen    = 225
	StatusClosingDataConnection = 226
	StatusPassiveMode           = 227
	StatusLongPassiveMode       = 228
	StatusExtendedPassiveMode   = 229
	StatusLoggedIn              = 230
	StatusLoggedOut             = 231
	StatusLogoutAck             = 232
	StatusRequestedFileActionOK = 250
	StatusPathCreated           = 257

	StatusUserOK             = 331
	StatusLoginNeedAccount   = 332
	StatusRequestFilePending = 350

	StatusNotAvailable             = 421
	StatusCanNotOpenDataConnection = 425
	StatusTransfertAborted         = 426
	StatusInvalidCredentials       = 430
	StatusHostUnavailable          = 434
	StatusFileActionIgnored        = 450
	StatusActionAborted            = 451
	Status452                      = 452

	StatusBadCommand              = 500
	StatusBadArguments            = 501
	StatusNotImplemented          = 502
	StatusBadSequence             = 503
	StatusNotImplementedParameter = 504
	StatusNotLoggedIn             = 530
	StatusStorNeedAccount         = 532
	StatusFileUnavailable         = 550
	StatusPageTypeUnknown         = 551
	StatusExceededStorage         = 552
	StatusBadFileName             = 553
)

var statusText = map[int]string{
	// 200
	StatusCommandOK:             "Command okay.",
	StatusCommandNotImplemented: "Command not implemented, superfluous at this site.",
	StatusSystem:                "System status, or system help reply.",
	StatusDirectory:             "Directory status.",
	StatusFile:                  "File status.",
	StatusHelp:                  "Help message.",
	StatusName:                  "",
	StatusReady:                 "Service ready for new user.",
	StatusClosing:               "Service closing control connection.",
	StatusDataConnectionOpen:    "Data connection open; no transfer in progress.",
	StatusClosingDataConnection: "Closing data connection. Requested file action successful.",
	StatusPassiveMode:           "Entering Passive Mode.",
	StatusLongPassiveMode:       "Entering Long Passive Mode.",
	StatusExtendedPassiveMode:   "Entering Extended Passive Mode.",
	StatusLoggedIn:              "User logged in, proceed.",
	StatusLoggedOut:             "User logged out; service terminated.",
	StatusLogoutAck:             "Logout command noted, will complete when transfer done.",
	StatusRequestedFileActionOK: "Requested file action okay, completed.",
	StatusPathCreated:           "Path created.",

	// 300
	StatusUserOK:             "User name okay, need password.",
	StatusLoginNeedAccount:   "Need account for login.",
	StatusRequestFilePending: "Requested file action pending further information.",

	// 400
	StatusNotAvailable:             "Service not available, closing control connection.",
	StatusCanNotOpenDataConnection: "Can't open data connection.",
	StatusTransfertAborted:         "Connection closed; transfer aborted.",
	StatusInvalidCredentials:       "Invalid username or password.",
	StatusHostUnavailable:          "Requested host unavailable.",
	StatusFileActionIgnored:        "Requested file action not taken.",
	StatusActionAborted:            "Requested action aborted. Local error in processing.",
	Status452:                      "Insufficient storage space in system.",

	// 500
	StatusBadCommand:              "Command unrecognized.",
	StatusBadArguments:            "Syntax error in parameters or arguments.",
	StatusNotImplemented:          "Command not implemented.",
	StatusBadSequence:             "Bad sequence of commands.",
	StatusNotImplementedParameter: "Command not implemented for that parameter.",
	StatusNotLoggedIn:             "Not logged in.",
	StatusStorNeedAccount:         "Need account for storing files.",
	StatusFileUnavailable:         "File unavailable.",
	StatusPageTypeUnknown:         "Page type unknown.",
	StatusExceededStorage:         "Exceeded storage allocation.",
	StatusBadFileName:             "File name not allowed.",
}

// StatusText returns a text for the FTP status code. It returns the empty string if the code is unknown.
func StatusText(code int) string {
	return statusText[code]
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"time"

	"github.com/cfpb/rhobot/config"
	"github.com/cfpb/rhobot/database"
	"github.com/cfpb/rhobot/gocd"
	"github.com/cfpb/rhobot/healthcheck"
	"github.com/cfpb/rhobot/report"
	"github.com/codegangsta/cli"
)

func main() {

	app := cli.NewApp()
	app.Name = "Rhobot"
	app.Usage = "Rhobot is a database development tool that uses DevOps best practices."
	app.EnableBashCompletion = true

	config := config.NewConfig()

	gocdHostFlag := cli.StringFlag{
		Name:  "host",
		Value: "",
		Usage: "host of the GoCD server",
	}
	reportFileFlag := cli.StringFlag{
		Name:  "report",
		Value: "",
		Usage: "path to the healthcheck report",
	}
	dburiFlag := cli.StringFlag{
		Name:  "dburi",
		Value: "",
		Usage: "database uri postgres://user:password@host:port/database",
	}

	app.Commands = []cli.Command{
		{
			Name:    "run",
			Aliases: []string{},
			Usage:   "healthchecks|pipeline|tbd",
			Subcommands: []cli.Command{
				{
					Name:  "healthchecks",
					Usage: "HEALTHCHECK_FILE [--dburi DATABASE_URI] [--report REPORT_FILE]",
					Flags: []cli.Flag{
						reportFileFlag,
						dburiFlag,
					},
					Action: func(c *cli.Context) {

						fmt.Println("ARGS: ", c.Args())

						if c.String("dburi") != "" {
							config.SetDBURI(c.String("dburi"))
						}
						fmt.Println("DB_URI: ", config.DBURI())

						var path string
						if c.Args().Get(0) != "" {
							path = c.Args().Get(0)
						} else {
							fmt.Println("You must provide the path to the healthcheck file.")
						}
						fmt.Println("PATH: ", path)

						healthChecks := healthcheck.ReadYamlFromFile(path)
						cxn := database.GetPGConnection(config.DBURI())
						results, _ := healthcheck.PreformHealthChecks(healthChecks, cxn)
						metadata := map[string]interface{}{
							"name":      healthChecks.Name,
							"db_name":   config.PgDatabase,
							"footer":    healthcheck.FooterHealthcheck,
							"timestamp": time.Now().UTC().String(),
						}

						var elements []report.Element
						for _, val := range results {
							elements = append(elements, val)
						}

						if c.String("report") != "" {
							prr := report.NewPongo2ReportRunnerFromString(healthcheck.TemplateHealthcheck)
							fhr := report.FileHandler{Filename: c.String("report")}
							rs := report.Set{Elements: elements, Metadata: metadata}
							reader, _ := prr.ReportReader(rs)
							_ = fhr.HandleReport(reader)
						}

					},
				},
				{
					Name:    "pipeline",
					Aliases: []string{},
					Usage:   "Interact with GoCD pipeline",
					Subcommands: []cli.Command{
						{
							Name:  "push",
							Usage: "PATH [PIPELINE_GROUP]",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								if len(c.Args()) > 0 {
									path := c.Args()[0]
									group := c.Args().Get(1)
									gocd.Push(config.GoCDURL(), path, group)
								} else {
									fmt.Println("PATH is required for push command.")
								}
							},
						},
						{
							Name:  "pull",
							Usage: "PATH",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								path := c.Args()[0]

								gocd.Pull(config.GoCDURL(), path)
							},
						},
						{
							Name:  "clone",
							Usage: "PIPELINE_NAME PATH",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								name := c.Args()[0]
								path := c.Args()[1]

								gocd.Clone(config.GoCDURL(), path, name)
							},
						},
					},
				},
			},
		},
	}

	app.Run(os.Args)
}
<commit_msg>cleaner cli argument variables<commit_after>package main

import (
	"fmt"
	"os"
	"time"

	"github.com/cfpb/rhobot/config"
	"github.com/cfpb/rhobot/database"
	"github.com/cfpb/rhobot/gocd"
	"github.com/cfpb/rhobot/healthcheck"
	"github.com/cfpb/rhobot/report"
	"github.com/codegangsta/cli"
)

func main() {

	app := cli.NewApp()
	app.Name = "Rhobot"
	app.Usage = "Rhobot is a database development tool that uses DevOps best practices."
	app.EnableBashCompletion = true

	config := config.NewConfig()

	gocdHostFlag := cli.StringFlag{
		Name:  "host",
		Value: "",
		Usage: "host of the GoCD server",
	}
	reportFileFlag := cli.StringFlag{
		Name:  "report",
		Value: "",
		Usage: "path to the healthcheck report",
	}
	dburiFlag := cli.StringFlag{
		Name:  "dburi",
		Value: "",
		Usage: "database uri postgres://user:password@host:port/database",
	}

	app.Commands = []cli.Command{
		{
			Name:    "run",
			Aliases: []string{},
			Usage:   "healthchecks|pipeline|tbd",
			Subcommands: []cli.Command{
				{
					Name:  "healthchecks",
					Usage: "HEALTHCHECK_FILE [--dburi DATABASE_URI] [--report REPORT_FILE]",
					Flags: []cli.Flag{
						reportFileFlag,
						dburiFlag,
					},
					Action: func(c *cli.Context) {

						// variables to be populated by cli args
						var healthcheckPath string
						var reportPath string

						if c.String("dburi") != "" {
							config.SetDBURI(c.String("dburi"))
						}

						if c.Args().Get(0) != "" {
							healthcheckPath = c.Args().Get(0)
						} else {
							fmt.Println("You must provide the path to the healthcheck file.")
						}

						if c.String("report") != "" {
							reportPath = c.String("report")
						}

						fmt.Println("DB_URI: ", config.DBURI())
						fmt.Println("PATH: ", healthcheckPath)

						healthChecks := healthcheck.ReadYamlFromFile(healthcheckPath)
						cxn := database.GetPGConnection(config.DBURI())
						results, _ := healthcheck.PreformHealthChecks(healthChecks, cxn)
						metadata := map[string]interface{}{
							"name":      healthChecks.Name,
							"db_name":   config.PgDatabase,
							"footer":    healthcheck.FooterHealthcheck,
							"timestamp": time.Now().UTC().String(),
						}

						var elements []report.Element
						for _, val := range results {
							elements = append(elements, val)
						}

						if reportPath != "" {
							prr := report.NewPongo2ReportRunnerFromString(healthcheck.TemplateHealthcheck)
							fhr := report.FileHandler{Filename: reportPath}
							rs := report.Set{Elements: elements, Metadata: metadata}
							reader, _ := prr.ReportReader(rs)
							_ = fhr.HandleReport(reader)
						}

					},
				},
				{
					Name:    "pipeline",
					Aliases: []string{},
					Usage:   "Interact with GoCD pipeline",
					Subcommands: []cli.Command{
						{
							Name:  "push",
							Usage: "PATH [PIPELINE_GROUP]",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								if len(c.Args()) > 0 {
									path := c.Args()[0]
									group := c.Args().Get(1)
									gocd.Push(config.GoCDURL(), path, group)
								} else {
									fmt.Println("PATH is required for push command.")
								}
							},
						},
						{
							Name:  "pull",
							Usage: "PATH",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								path := c.Args()[0]

								gocd.Pull(config.GoCDURL(), path)
							},
						},
						{
							Name:  "clone",
							Usage: "PIPELINE_NAME PATH",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								name := c.Args()[0]
								path := c.Args()[1]

								gocd.Clone(config.GoCDURL(), path, name)
							},
						},
					},
				},
			},
		},
	}

	app.Run(os.Args)
}
<|endoftext|>