<commit_before>package build

import (
	"bytes"
	"encoding/gob"
	"errors"
	"strconv"

	"github.com/midlang/mid/src/mid/ast"
	"github.com/midlang/mid/src/mid/lexer"
	"github.com/mkideal/log"
)

var (
	ErrAmbiguousNames = errors.New("ambiguous names")
)

func init() {
	gob.Register(&Field{})
	gob.Register(&Bean{})
	gob.Register(&File{})
	gob.Register(&Package{})

	gob.Register(ExprBase{})
	gob.Register(TypeBase{})
	gob.Register(&BasicLit{})
	gob.Register(&BasicType{})
	gob.Register(&ArrayType{})
	gob.Register(&MapType{})
	gob.Register(&VectorType{})
	gob.Register(&GenDecl{})
	gob.Register(&ImportSpec{})
	gob.Register(&ConstSpec{})
	gob.Register(&StructType{})
	gob.Register(&FuncType{})
}

// helper functions
func BuildDoc(doc *ast.CommentGroup) string {
	if doc == nil {
		return ""
	}
	text := doc.Text()
	if text != "" {
		text += "\n"
	}
	return text
}

func BuildComment(comment *ast.CommentGroup) string {
	if comment == nil {
		return ""
	}
	text := comment.Text()
	return text
}

func BuildTag(tag *ast.BasicLit) Tag {
	if tag == nil {
		return ""
	}
	return Tag(tag.Value)
}

func BuildIdent(ident *ast.Ident) string {
	if ident == nil {
		return ""
	}
	return ident.Name
}

func BuildIdentList(idents []*ast.Ident) []string {
	if len(idents) == 0 {
		return []string{}
	}
	strs := make([]string, 0, len(idents))
	for _, ident := range idents {
		strs = append(strs, BuildIdent(ident))
	}
	return strs
}

type Tag string

func (tag Tag) Get(key string) string {
	value, _ := tag.Lookup(key)
	return value
}

func (tag *Tag) Set(key, value string) {
	pairs, _, index := tag.parse(key)
	if index >= 0 {
		pairs[index][1] = value
	} else {
		pairs = append(pairs, tagpair{key, value})
	}
	*tag = Tag(tag.format(pairs))
}

func (tag Tag) Lookup(key string) (value string, ok bool) {
	_, value, index := tag.parse(key)
	return value, index >= 0
}

type tagpair [2]string

func (tag Tag) format(pairs []tagpair) string {
	var buf bytes.Buffer
	for i, pair := range pairs {
		if i > 0 {
			buf.WriteByte(' ')
		}
		buf.WriteString(pair[0])
		buf.WriteByte(':')
		buf.WriteByte('"')
		buf.WriteString(pair[1])
		buf.WriteByte('"')
	}
	return buf.String()
}

func (tag Tag) parse(key string) (pairs []tagpair, value string, index int) {
	pairs = make([]tagpair, 0)
	index = -1
	count := 0
	for tag != "" {
		// Skip leading space.
		i := 0
		for i < len(tag) && tag[i] == ' ' {
			i++
		}
		tag = tag[i:]
		if tag == "" {
			break
		}
		i = 0
		for i < len(tag) && tag[i] > ' ' && tag[i] != ':' && tag[i] != '"' && tag[i] != 0x7f {
			i++
		}
		if i == 0 || i+1 >= len(tag) || tag[i] != ':' || tag[i+1] != '"' {
			break
		}
		name := string(tag[:i])
		tag = tag[i+1:]

		// Scan quoted string to find value.
		i = 1
		for i < len(tag) && tag[i] != '"' {
			if tag[i] == '\\' {
				i++
			}
			i++
		}
		if i >= len(tag) {
			break
		}
		qvalue := string(tag[:i+1])
		tag = tag[i+1:]

		v, err := strconv.Unquote(qvalue)
		if err != nil {
			continue
		}
		if key == name && key != "" && index == -1 {
			value = v
			index = len(pairs)
		}
		pairs = append(pairs, tagpair{name, v})
		count++
	}
	return
}

type Field struct {
	Doc     string
	Options []string
	Type    Type
	Names   []string
	Default Expr
	Tag     Tag
	Comment string
}

func (field Field) Name() (string, error) {
	if len(field.Names) == 0 {
		return "", nil
	}
	if len(field.Names) == 1 {
		return field.Names[0], nil
	}
	return "", ErrAmbiguousNames
}

func (field Field) Value() string {
	switch e := field.Default.(type) {
	case *BasicLit:
		return e.Value
	}
	panic("unsupported expr")
}

func (field Field) IsExtended() bool {
	return len(field.Names) == 0
}

func (field Field) GetTag(key string) string {
	return field.Tag.Get(key)
}

func (field Field) HasTag(key string) bool {
	_, ok := field.Tag.Lookup(key)
	return ok
}

func (field *Field) AddTag(key, value string) {
	field.Tag.Set(key, value)
}

func BuildField(field *ast.Field) *Field {
	out := &Field{
		Doc:     BuildDoc(field.Doc),
		Options: BuildIdentList(field.Options),
		Type:    BuildType(field.Type),
		Names:   BuildIdentList(field.Names),
		Default: BuildExpr(field.Default),
		Tag:     BuildTag(field.Tag),
		Comment: BuildComment(field.Comment),
	}
	log.Trace("BuildField: field=%v", out.Names)
	return out
}

func BuildFieldList(fields *ast.FieldList) []*Field {
	if fields == nil || len(fields.List) == 0 {
		return []*Field{}
	}
	list := make([]*Field, 0, len(fields.List))
	for _, field := range fields.List {
		list = append(list, BuildField(field))
	}
	return list
}

type Expr interface {
	ExprNode()
}

type ExprBase struct{}

func (ExprBase) ExprNode() {}
func (Ident) ExprNode()    {}
func (BasicLit) ExprNode() {}

type Ident string

type BasicLit struct {
	Kind  lexer.Token
	Value string
}

func BuildBasicLit(lit *ast.BasicLit) *BasicLit {
	return &BasicLit{
		Kind:  lit.Tok,
		Value: lit.Value,
	}
}

func BuildExpr(expr ast.Expr) Expr {
	if typ, ok := expr.(ast.Type); ok {
		return BuildType(typ)
	}
	if ident, ok := expr.(*ast.Ident); ok {
		return Ident(BuildIdent(ident))
	}
	if lit, ok := expr.(*ast.BasicLit); ok {
		return BuildBasicLit(lit)
	}
	// TODO: alert error
	return &ExprBase{}
}

type Type interface {
	Expr
	TypeNode()
}

type TypeBase struct {
	ExprBase
}

func (TypeBase) TypeNode() {}

func BuildType(typ ast.Type) Type {
	switch t := typ.(type) {
	case *ast.BasicType:
		return &BasicType{Name: BuildIdent(t.Name)}
	case *ast.StructType:
		return BuildStruct(t)
	case *ast.MapType:
		return BuildMap(t)
	case *ast.ArrayType:
		return BuildArray(t)
	case *ast.VectorType:
		return BuildVector(t)
	case *ast.FuncType:
		return BuildFunc(t)
	default:
		return &TypeBase{}
	}
}

type BasicType struct {
	TypeBase
	Name string
}

type ArrayType struct {
	TypeBase
	T    Type
	Size Expr
}

func BuildArray(t *ast.ArrayType) *ArrayType {
	return &ArrayType{
		T:    BuildType(t.T),
		Size: BuildExpr(t.Size),
	}
}

type MapType struct {
	TypeBase
	K Type
	V Type
}

func BuildMap(t *ast.MapType) *MapType {
	return &MapType{
		K: BuildType(t.K),
		V: BuildType(t.V),
	}
}

type VectorType struct {
	TypeBase
	T Type
}

func BuildVector(t *ast.VectorType) *VectorType {
	return &VectorType{
		T: BuildType(t.T),
	}
}

type StructType struct {
	TypeBase
	Package string
	Name    string
}

func BuildStruct(t *ast.StructType) *StructType {
	return &StructType{
		Package: BuildIdent(t.Package),
		Name:    BuildIdent(t.Name),
	}
}

type FuncType struct {
	TypeBase
	Params []*Field
	Result Type
}

func BuildFunc(t *ast.FuncType) *FuncType {
	return &FuncType{
		Params: BuildFieldList(t.Params),
		Result: BuildType(t.Result),
	}
}

type Bean struct {
	Kind    string
	Doc     string
	Name    string
	Fields  []*Field
	Comment string
}

func BuildBean(bean *ast.BeanDecl) *Bean {
	log.Trace("BuildBean: kind=%s, name=%s, len(fields)=%d", bean.Kind, bean.Name.Name, len(bean.Fields.List))
	b := &Bean{
		Kind:   bean.Kind,
		Doc:    BuildDoc(bean.Doc),
		Name:   BuildIdent(bean.Name),
		Fields: BuildFieldList(bean.Fields),
	}
	log.Trace("BuildBean: %v", bean)
	return b
}

func (bean Bean) Extends() []Type {
	var extends []Type
	for _, f := range bean.Fields {
		if len(f.Names) == 0 {
			extends = append(extends, f.Type)
		}
	}
	return extends
}

func (bean Bean) Field(i int) *Field {
	if i >= len(bean.Fields) || i < 0 {
		return nil
	}
	return bean.Fields[i]
}

func (bean Bean) NumField() int { return len(bean.Fields) }

type ImportSpec struct {
	Doc     string
	Name    string
	Package string
	Comment string
}

func BuildImportSpec(spec *ast.ImportSpec) *ImportSpec {
	return &ImportSpec{
		Doc:     BuildDoc(spec.Doc),
		Name:    BuildIdent(spec.Name),
		Package: spec.Package.Value,
		Comment: BuildDoc(spec.Comment),
	}
}

type ConstSpec struct {
	Doc     string
	Name    string
	Value   Expr
	Comment string
}

func (c ConstSpec) ValueString() string {
	switch e := c.Value.(type) {
	case *BasicLit:
		return e.Value
	}
	panic("unsupported expr")
}

func BuildConstSpec(spec *ast.ConstSpec) *ConstSpec {
	return &ConstSpec{
		Doc:     BuildDoc(spec.Doc),
		Name:    BuildIdent(spec.Name),
		Value:   BuildExpr(spec.Value),
		Comment: BuildComment(spec.Comment),
	}
}

type GenDecl struct {
	Doc     string
	Imports []*ImportSpec
	Consts  []*ConstSpec
}

func BuildGenDecl(decl *ast.GenDecl) *GenDecl {
	d := &GenDecl{
		Doc: BuildDoc(decl.Doc),
	}
	for _, spec := range decl.Specs {
		switch s := spec.(type) {
		case *ast.ImportSpec:
			d.Imports = append(d.Imports, BuildImportSpec(s))
		case *ast.ConstSpec:
			d.Consts = append(d.Consts, BuildConstSpec(s))
		default:
			//TODO: alert error
		}
	}
	return d
}

type File struct {
	Filename   string
	Doc        string
	Package    string
	Beans      []*Bean
	Decls      []*GenDecl
	Unresolved []string
}

func BuildFile(file *ast.File) *File {
	f := &File{
		Filename:   file.Filename,
		Doc:        BuildDoc(file.Doc),
		Package:    BuildIdent(file.Name),
		Unresolved: BuildIdentList(file.Unresolved),
	}
	for _, decl := range file.Decls {
		switch d := decl.(type) {
		case *ast.BeanDecl:
			f.Beans = append(f.Beans, BuildBean(d))
		case *ast.GenDecl:
			f.Decls = append(f.Decls, BuildGenDecl(d))
		}
	}
	return f
}

type Package struct {
	Name    string
	Imports map[string]string
	Files   map[string]*File
}

func BuildPackage(pkg *ast.Package) *Package {
	p := &Package{
		Name:    pkg.Name,
		Imports: make(map[string]string),
		Files:   make(map[string]*File),
	}
	if pkg.Imports != nil {
		for id, name := range pkg.Imports {
			p.Imports[id] = name.Name
		}
	}
	if pkg.Files != nil {
		for name, file := range pkg.Files {
			p.Files[name] = BuildFile(file)
		}
	}
	return p
}
<commit_msg>add IsArray,IsVector and IsMap for build.Type<commit_after>package build

import (
	"bytes"
	"encoding/gob"
	"errors"
	"strconv"

	"github.com/midlang/mid/src/mid/ast"
	"github.com/midlang/mid/src/mid/lexer"
	"github.com/mkideal/log"
)

var (
	ErrAmbiguousNames = errors.New("ambiguous names")
)

func init() {
	gob.Register(&Field{})
	gob.Register(&Bean{})
	gob.Register(&File{})
	gob.Register(&Package{})

	gob.Register(ExprBase{})
	gob.Register(TypeBase{})
	gob.Register(&BasicLit{})
	gob.Register(&BasicType{})
	gob.Register(&ArrayType{})
	gob.Register(&MapType{})
	gob.Register(&VectorType{})
	gob.Register(&GenDecl{})
	gob.Register(&ImportSpec{})
	gob.Register(&ConstSpec{})
	gob.Register(&StructType{})
	gob.Register(&FuncType{})
}

// helper functions
func BuildDoc(doc *ast.CommentGroup) string {
	if doc == nil {
		return ""
	}
	text := doc.Text()
	if text != "" {
		text += "\n"
	}
	return text
}

func BuildComment(comment *ast.CommentGroup) string {
	if comment == nil {
		return ""
	}
	text := comment.Text()
	return text
}

func BuildTag(tag *ast.BasicLit) Tag {
	if tag == nil {
		return ""
	}
	return Tag(tag.Value)
}

func BuildIdent(ident *ast.Ident) string {
	if ident == nil {
		return ""
	}
	return ident.Name
}

func BuildIdentList(idents []*ast.Ident) []string {
	if len(idents) == 0 {
		return []string{}
	}
	strs := make([]string, 0, len(idents))
	for _, ident := range idents {
		strs = append(strs, BuildIdent(ident))
	}
	return strs
}

type Tag string

func (tag Tag) Get(key string) string {
	value, _ := tag.Lookup(key)
	return value
}

func (tag *Tag) Set(key, value string) {
	pairs, _, index := tag.parse(key)
	if index >= 0 {
		pairs[index][1] = value
	} else {
		pairs = append(pairs, tagpair{key, value})
	}
	*tag = Tag(tag.format(pairs))
}

func (tag Tag) Lookup(key string) (value string, ok bool) {
	_, value, index := tag.parse(key)
	return value, index >= 0
}

type tagpair [2]string

func (tag Tag) format(pairs []tagpair) string {
	var buf bytes.Buffer
	for i, pair := range pairs {
		if i > 0 {
			buf.WriteByte(' ')
		}
		buf.WriteString(pair[0])
		buf.WriteByte(':')
		buf.WriteByte('"')
		buf.WriteString(pair[1])
		buf.WriteByte('"')
	}
	return buf.String()
}

func (tag Tag) parse(key string) (pairs []tagpair, value string, index int) {
	pairs = make([]tagpair, 0)
	index = -1
	count := 0
	for tag != "" {
		// Skip leading space.
		i := 0
		for i < len(tag) && tag[i] == ' ' {
			i++
		}
		tag = tag[i:]
		if tag == "" {
			break
		}
		i = 0
		for i < len(tag) && tag[i] > ' ' && tag[i] != ':' && tag[i] != '"' && tag[i] != 0x7f {
			i++
		}
		if i == 0 || i+1 >= len(tag) || tag[i] != ':' || tag[i+1] != '"' {
			break
		}
		name := string(tag[:i])
		tag = tag[i+1:]

		// Scan quoted string to find value.
		i = 1
		for i < len(tag) && tag[i] != '"' {
			if tag[i] == '\\' {
				i++
			}
			i++
		}
		if i >= len(tag) {
			break
		}
		qvalue := string(tag[:i+1])
		tag = tag[i+1:]

		v, err := strconv.Unquote(qvalue)
		if err != nil {
			continue
		}
		if key == name && key != "" && index == -1 {
			value = v
			index = len(pairs)
		}
		pairs = append(pairs, tagpair{name, v})
		count++
	}
	return
}

type Field struct {
	Doc     string
	Options []string
	Type    Type
	Names   []string
	Default Expr
	Tag     Tag
	Comment string
}

func (field Field) Name() (string, error) {
	if len(field.Names) == 0 {
		return "", nil
	}
	if len(field.Names) == 1 {
		return field.Names[0], nil
	}
	return "", ErrAmbiguousNames
}

func (field Field) Value() string {
	switch e := field.Default.(type) {
	case *BasicLit:
		return e.Value
	}
	panic("unsupported expr")
}

func (field Field) IsExtended() bool {
	return len(field.Names) == 0
}

func (field Field) GetTag(key string) string {
	return field.Tag.Get(key)
}

func (field Field) HasTag(key string) bool {
	_, ok := field.Tag.Lookup(key)
	return ok
}

func (field *Field) AddTag(key, value string) {
	field.Tag.Set(key, value)
}

func BuildField(field *ast.Field) *Field {
	out := &Field{
		Doc:     BuildDoc(field.Doc),
		Options: BuildIdentList(field.Options),
		Type:    BuildType(field.Type),
		Names:   BuildIdentList(field.Names),
		Default: BuildExpr(field.Default),
		Tag:     BuildTag(field.Tag),
		Comment: BuildComment(field.Comment),
	}
	log.Trace("BuildField: field=%v", out.Names)
	return out
}

func BuildFieldList(fields *ast.FieldList) []*Field {
	if fields == nil || len(fields.List) == 0 {
		return []*Field{}
	}
	list := make([]*Field, 0, len(fields.List))
	for _, field := range fields.List {
		list = append(list, BuildField(field))
	}
	return list
}

type Expr interface {
	ExprNode()
}

type ExprBase struct{}

func (ExprBase) ExprNode() {}
func (Ident) ExprNode()    {}
func (BasicLit) ExprNode() {}

type Ident string

type BasicLit struct {
	Kind  lexer.Token
	Value string
}

func BuildBasicLit(lit *ast.BasicLit) *BasicLit {
	return &BasicLit{
		Kind:  lit.Tok,
		Value: lit.Value,
	}
}

func BuildExpr(expr ast.Expr) Expr {
	if typ, ok := expr.(ast.Type); ok {
		return BuildType(typ)
	}
	if ident, ok := expr.(*ast.Ident); ok {
		return Ident(BuildIdent(ident))
	}
	if lit, ok := expr.(*ast.BasicLit); ok {
		return BuildBasicLit(lit)
	}
	// TODO: alert error
	return &ExprBase{}
}

type Type interface {
	Expr
	TypeNode()
	IsArray() bool
	IsVector() bool
	IsMap() bool
	IsStruct() bool
}

type TypeBase struct {
	ExprBase
}

func (TypeBase) TypeNode()      {}
func (TypeBase) IsArray() bool  { return false }
func (TypeBase) IsVector() bool { return false }
func (TypeBase) IsMap() bool    { return false }
func (TypeBase) IsStruct() bool { return false }

func BuildType(typ ast.Type) Type {
	switch t := typ.(type) {
	case *ast.BasicType:
		return &BasicType{Name: BuildIdent(t.Name)}
	case *ast.StructType:
		return BuildStruct(t)
	case *ast.MapType:
		return BuildMap(t)
	case *ast.ArrayType:
		return BuildArray(t)
	case *ast.VectorType:
		return BuildVector(t)
	case *ast.FuncType:
		return BuildFunc(t)
	default:
		return &TypeBase{}
	}
}

type BasicType struct {
	TypeBase
	Name string
}

type ArrayType struct {
	TypeBase
	T    Type
	Size Expr
}

func (ArrayType) IsArray() bool { return true }

func BuildArray(t *ast.ArrayType) *ArrayType {
	return &ArrayType{
		T:    BuildType(t.T),
		Size: BuildExpr(t.Size),
	}
}

type MapType struct {
	TypeBase
	K Type
	V Type
}

func (MapType) IsMap() bool { return true }

func BuildMap(t *ast.MapType) *MapType {
	return &MapType{
		K: BuildType(t.K),
		V: BuildType(t.V),
	}
}

type VectorType struct {
	TypeBase
	T Type
}

func (VectorType) IsVector() bool { return true }

func BuildVector(t *ast.VectorType) *VectorType {
	return &VectorType{
		T: BuildType(t.T),
	}
}

type StructType struct {
	TypeBase
	Package string
	Name    string
}

func (StructType) IsStruct() bool { return true }

func BuildStruct(t *ast.StructType) *StructType {
	return &StructType{
		Package: BuildIdent(t.Package),
		Name:    BuildIdent(t.Name),
	}
}

type FuncType struct {
	TypeBase
	Params []*Field
	Result Type
}

func BuildFunc(t *ast.FuncType) *FuncType {
	return &FuncType{
		Params: BuildFieldList(t.Params),
		Result: BuildType(t.Result),
	}
}

type Bean struct {
	Kind    string
	Doc     string
	Name    string
	Fields  []*Field
	Comment string
}

func BuildBean(bean *ast.BeanDecl) *Bean {
	log.Trace("BuildBean: kind=%s, name=%s, len(fields)=%d", bean.Kind, bean.Name.Name, len(bean.Fields.List))
	b := &Bean{
		Kind:   bean.Kind,
		Doc:    BuildDoc(bean.Doc),
		Name:   BuildIdent(bean.Name),
		Fields: BuildFieldList(bean.Fields),
	}
	log.Trace("BuildBean: %v", bean)
	return b
}

func (bean Bean) Extends() []Type {
	var extends []Type
	for _, f := range bean.Fields {
		if len(f.Names) == 0 {
			extends = append(extends, f.Type)
		}
	}
	return extends
}

func (bean Bean) Field(i int) *Field {
	if i >= len(bean.Fields) || i < 0 {
		return nil
	}
	return bean.Fields[i]
}

func (bean Bean) NumField() int { return len(bean.Fields) }

type ImportSpec struct {
	Doc     string
	Name    string
	Package string
	Comment string
}

func BuildImportSpec(spec *ast.ImportSpec) *ImportSpec {
	return &ImportSpec{
		Doc:     BuildDoc(spec.Doc),
		Name:    BuildIdent(spec.Name),
		Package: spec.Package.Value,
		Comment: BuildDoc(spec.Comment),
	}
}

type ConstSpec struct {
	Doc     string
	Name    string
	Value   Expr
	Comment string
}

func (c ConstSpec) ValueString() string {
	switch e := c.Value.(type) {
	case *BasicLit:
		return e.Value
	}
	panic("unsupported expr")
}

func BuildConstSpec(spec *ast.ConstSpec) *ConstSpec {
	return &ConstSpec{
		Doc:     BuildDoc(spec.Doc),
		Name:    BuildIdent(spec.Name),
		Value:   BuildExpr(spec.Value),
		Comment: BuildComment(spec.Comment),
	}
}

type GenDecl struct {
	Doc     string
	Imports []*ImportSpec
	Consts  []*ConstSpec
}

func BuildGenDecl(decl *ast.GenDecl) *GenDecl {
	d := &GenDecl{
		Doc: BuildDoc(decl.Doc),
	}
	for _, spec := range decl.Specs {
		switch s := spec.(type) {
		case *ast.ImportSpec:
			d.Imports = append(d.Imports, BuildImportSpec(s))
		case *ast.ConstSpec:
			d.Consts = append(d.Consts, BuildConstSpec(s))
		default:
			//TODO: alert error
		}
	}
	return d
}

type File struct {
	Filename   string
	Doc        string
	Package    string
	Beans      []*Bean
	Decls      []*GenDecl
	Unresolved []string
}

func BuildFile(file *ast.File) *File {
	f := &File{
		Filename:   file.Filename,
		Doc:        BuildDoc(file.Doc),
		Package:    BuildIdent(file.Name),
		Unresolved: BuildIdentList(file.Unresolved),
	}
	for _, decl := range file.Decls {
		switch d := decl.(type) {
		case *ast.BeanDecl:
			f.Beans = append(f.Beans, BuildBean(d))
		case *ast.GenDecl:
			f.Decls = append(f.Decls, BuildGenDecl(d))
		}
	}
	return f
}

type Package struct {
	Name    string
	Imports map[string]string
	Files   map[string]*File
}

func BuildPackage(pkg *ast.Package) *Package {
	p := &Package{
		Name:    pkg.Name,
		Imports: make(map[string]string),
		Files:   make(map[string]*File),
	}
	if pkg.Imports != nil {
		for id, name := range pkg.Imports {
			p.Imports[id] = name.Name
		}
	}
	if pkg.Files != nil {
		for name, file := range pkg.Files {
			p.Files[name] = BuildFile(file)
		}
	}
	return p
}
<|endoftext|>


<commit_before>package triggerfail

import (
	"io"
	"os/exec"
	"strings"
)

// Run a command, checking if triggers were found in it's output.
// It returns a list of triggers found in the stdout and stderr of the running command.
// This command takes the following arguments:
//   *exec.Cmd - a pointer to an exec.Cmd (usually created by the exec.Command func). It can be configured as usual with the exception of stdin and stdout which should be left as is.
//   []string  - A list of trigger strings. The command will be failed if it's stdout or stderr contains these strings.
//   io.Writer - After consuming the stdout of the command to check for triggers, we pass along the results to this writer. Usually it would be set to os.Stdout or os.DevNull
//   io.Writer - After consuming the stderr of the command to check for triggers, we pass along the results to this writer. Usually it would be set to os.Stderr or os.DevNull
//   bool      - Abort the running commmand if a trigger word is found. If set to false the command will be allowed to run to completion, even with triggers.
func RunCommand(cmd *exec.Cmd, triggers []string, stdout io.Writer, stderr io.Writer, abort bool) ([]string, error) {
	found := make([]string, 0)

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return found, err
	}
	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return found, err
	}
	err = cmd.Start()
	if err != nil {
		return found, err
	}

	// Collect stdout and stderr
	stdoutbuff := make([]byte, 1024)
	stderrbuff := make([]byte, 1024)

	cmd.Run()

	for {
		//@@TODO: Move this to buffio.Scanner and two seperate co-running go-routines

		done := false

		// stdout
		n, err := stdoutPipe.Read(stdoutbuff)
		if err != nil {
			if err == io.EOF {
				done = true
			} else {
				return found, err
			}
		}
		stdout.Write(stdoutbuff[:n])

		// stderr
		n, err = stderrPipe.Read(stderrbuff)
		if err != nil {
			if err == io.EOF {
				done = true
			} else {
				return found, err
			}
		}
		stderr.Write(stderrbuff[:n])

		// If stdout contains a trigger, log it and possibly abort
		for _, trigger := range triggers {
			if strings.Contains(string(stdoutbuff), trigger) {
				found = append(found, trigger) //@@TODO: mutex
				if abort {
					if !done {
						cmd.Process.Kill() // abort the running command
					}
					done = true
				}
			}
		}

		// If stderr contains a trigger, log it and possibly abort
		for _, trigger := range triggers {
			if strings.Contains(string(stderrbuff), trigger) {
				found = append(found, trigger) //@@TODO: mutex
				if abort {
					if !done {
						cmd.Process.Kill() // abort the running command
					}
					done = true
				}
			}
		}

		if done {
			break
		}
	}

	err = cmd.Wait()
	if err != nil { // If we have an error abort everything
		if err.Error() == "signal: killed" { // We killed the process intentionally, don't report the error
			return found, nil
		} else {
			return found, err
		}
	}

	// Command complete
	return found, nil
}
<commit_msg>inline documentation<commit_after>package triggerfail

import (
	"io"
	"os/exec"
	"strings"
)

// Run a command, checking if triggers were found in it's output.
// It returns a list of triggers found in the stdout and stderr of the running command.
// This command takes the following arguments:
//   cmd      *exec.Cmd - a pointer to an exec.Cmd (usually created by the exec.Command func). It can be configured as usual with the exception of stdin and stdout which should be left as is.
//   triggers []string  - A list of trigger strings. The command will be failed if it's stdout or stderr contains these strings.
//   stdout   io.Writer - After consuming the stdout of the command to check for triggers, we pass along the results to this writer. Usually it would be set to os.Stdout or os.DevNull
//   stderr   io.Writer - After consuming the stderr of the command to check for triggers, we pass along the results to this writer. Usually it would be set to os.Stderr or os.DevNull
//   abort    bool      - Abort the running commmand if a trigger word is found. If set to false the command will be allowed to run to completion, even with triggers.
func RunCommand(cmd *exec.Cmd, triggers []string, stdout io.Writer, stderr io.Writer, abort bool) ([]string, error) {
	found := make([]string, 0)

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return found, err
	}
	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return found, err
	}
	err = cmd.Start()
	if err != nil {
		return found, err
	}

	// Collect stdout and stderr
	stdoutbuff := make([]byte, 1024)
	stderrbuff := make([]byte, 1024)

	cmd.Run()

	for {
		//@@TODO: Move this to buffio.Scanner and two seperate co-running go-routines

		done := false

		// stdout
		n, err := stdoutPipe.Read(stdoutbuff)
		if err != nil {
			if err == io.EOF {
				done = true
			} else {
				return found, err
			}
		}
		stdout.Write(stdoutbuff[:n])

		// stderr
		n, err = stderrPipe.Read(stderrbuff)
		if err != nil {
			if err == io.EOF {
				done = true
			} else {
				return found, err
			}
		}
		stderr.Write(stderrbuff[:n])

		// If stdout contains a trigger, log it and possibly abort
		for _, trigger := range triggers {
			if strings.Contains(string(stdoutbuff), trigger) {
				found = append(found, trigger) //@@TODO: mutex
				if abort {
					if !done {
						cmd.Process.Kill() // abort the running command
					}
					done = true
				}
			}
		}

		// If stderr contains a trigger, log it and possibly abort
		for _, trigger := range triggers {
			if strings.Contains(string(stderrbuff), trigger) {
				found = append(found, trigger) //@@TODO: mutex
				if abort {
					if !done {
						cmd.Process.Kill() // abort the running command
					}
					done = true
				}
			}
		}

		if done {
			break
		}
	}

	err = cmd.Wait()
	if err != nil { // If we have an error abort everything
		if err.Error() == "signal: killed" { // We killed the process intentionally, don't report the error
			return found, nil
		} else {
			return found, err
		}
	}

	// Command complete
	return found, nil
}
<|endoftext|>


<commit_before>package goat

import (
	"time"
)

type IoWorker struct {
	myLookup MapStorage
	db       *MapDb
}

// looks up where in the dbStor the requested data is stored
func (worker IoWorker) Lookup(id string) interface{} {
	if val, ok := worker.myLookup[id]; ok {
		return val
	} else {
		var err ErrorRes
		err.ErrLocation = "Lookup"
		err.Error = "id not found in map"
		err.Time = time.Now().Unix()
		var res Response
		res.Data = err
		res.Id = id
		Static.ErrChan <- res
		return nil
	}

}

// replaces the pointer in the lookup with a new pointer to the new data
func (worker IoWorker) UpdateLookup(id string, update interface{}) {
	worker.db.MapLookup[id] = &update
}

// Append new pointer to the end of a relation or create a new one if relation is not found
func (worker IoWorker) AppendRelation(id string, data *interface{}) {

	if hold, ok := worker.myLookup[id].(Relation); ok {
		hold.Index = append(hold.Index, data)
		worker.myLookup[id] = hold
	} else {
		n := new(Relation)
		n.Index[0] = data
		worker.myLookup[id] = n
	}
}

// Removes a given item from the index of a relation
func (worker IoWorker) RemoveRelation(id string, data *interface{}) {
	if hold, ok := worker.myLookup[id].(Relation); ok {
		for i := 0; i < len(hold.Index); i++ {
			if hold.Index[i] == data {
				hold.Index = append(hold.Index[:i], hold.Index[i+1:]...)
			}
		}
	} else {
		var err ErrorRes
		err.ErrLocation = "RemoveRelation"
		err.Error = "id not found in map"
		err.Time = time.Now().Unix()
		var res Response
		res.Data = err
		res.Id = id
		Static.ErrChan <- res
	}
}

// remove a pointer from lookup
func (worker IoWorker) RemoveLookup(id string) {
	worker.myLookup[id] = nil
}

// replace the data for a given id where ever it might be(which is identified by the Lookup function)
func (worker IoWorker) UpdateStor(id string, data interface{}, location map[string]interface{}) {
	location[id] = data
}
<commit_msg>added error handling to UpdateLookup<commit_after>package goat

import (
	"time"
)

type IoWorker struct {
	myLookup MapStorage
	db       *MapDb
}

// looks up where in the dbStor the requested data is stored
func (worker IoWorker) Lookup(id string) interface{} {
	if val, ok := worker.myLookup[id]; ok {
		return val
	} else {
		var err ErrorRes
		err.ErrLocation = "Lookup"
		err.Error = "id not found in map"
		err.Time = time.Now().Unix()
		var res Response
		res.Data = err
		res.Id = id
		Static.ErrChan <- res
		return nil
	}

}

// replaces the pointer in the lookup with a new pointer to the new data
func (worker IoWorker) UpdateLookup(id string, update interface{}) {
	if _, ok := worker.db.MapLookup[id]; ok {
		worker.db.MapLookup[id] = &update
	} else {
		var err ErrorRes
		err.ErrLocation = "UpdateLookup"
		err.Error = "id not found in map"
		err.Time = time.Now().Unix()
		var res Response
		res.Data = err
		res.Id = id
		Static.ErrChan <- res
	}

}

// Append new pointer to the end of a relation or create a new one if relation is not found
func (worker IoWorker) AppendRelation(id string, data *interface{}) {

	if hold, ok := worker.myLookup[id].(Relation); ok {
		hold.Index = append(hold.Index, data)
		worker.myLookup[id] = hold
	} else {
		n := new(Relation)
		n.Index[0] = data
		worker.myLookup[id] = n
	}
}

// Removes a given item from the index of a relation
func (worker IoWorker) RemoveRelation(id string, data *interface{}) {
	if hold, ok := worker.myLookup[id].(Relation); ok {
		for i := 0; i < len(hold.Index); i++ {
			if hold.Index[i] == data {
				hold.Index = append(hold.Index[:i], hold.Index[i+1:]...)
			}
		}
	} else {
		var err ErrorRes
		err.ErrLocation = "RemoveRelation"
		err.Error = "id not found in map"
		err.Time = time.Now().Unix()
		var res Response
		res.Data = err
		res.Id = id
		Static.ErrChan <- res
	}
}

// remove a pointer from lookup
func (worker IoWorker) RemoveLookup(id string) {
	worker.myLookup[id] = nil
}

// replace the data for a given id where ever it might be(which is identified by the Lookup function)
func (worker IoWorker) UpdateStor(id string, data interface{}, location map[string]interface{}) {
	location[id] = data
}
<|endoftext|>


<commit_before>package main

import (
		"fmt"
		"github.com/bs_radix"
)

func main() {
	fmt.Println("Starting Main")
	t := bs_radix.New()
	fmt.Println("t is ", t)
	t.Insert(":445L/PG", "445LPG")
	t.Insert(":445L/PG/RegDoc", "445LPG")
	t.Insert(":445L/PG/SampleDoc", "bob")
	t.Insert(":445L/PG/SampleCode", "bobby")
	val, found := t.Get(":445L/PG/SampleCode")
	if found {
		fmt.Println("445L/PG/SampleCode was found. Value: ", val)
	}
	val, found = t.Get(":445L/PG/SampleDoc")
	if found {
		fmt.Println("445L/PG/SampleDoc was found. Value: ", val)
	}
	//look for /Sample. Should return not found
	val, found = t.Get(":445L/PG/Sample")
	if found {
		fmt.Println("445L/PG/Sample was found. Value: ", val)
	}
	//now insert /Sample and check again
	t.Insert(":445L/PG/Sample", "i'm here")
	val, found = t.Get(":445L/PG/Sample")
	if found {
		fmt.Println("445L/PG/Sample was found. Value: ", val)
	}

}<commit_msg>remove radixtest<commit_after><|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/jamiealquiza/tachymeter"
	"gopkg.in/Shopify/sarama.v1"
)

type config struct {
	brokers          []string
	topic            string
	msgSize          int
	msgRate          uint64
	batchSize        int
	compression      sarama.CompressionCodec
	compressionName  string
	requiredAcks     sarama.RequiredAcks
	requiredAcksName string
	workers          int
	writersPerWorker int
	noop             bool
}

var (
	Config = &config{}

	// Character selection for random messages.
	chars = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$^&*(){}][:<>.")

	// Counters / misc.
	sentCnt uint64
	errCnt  uint64
)

func init() {
	flag.StringVar(&Config.topic, "topic", "sangrenel", "Kafka topic to produce to")
	flag.IntVar(&Config.msgSize, "message-size", 300, "Message size (bytes)")
	flag.Uint64Var(&Config.msgRate, "produce-rate", 100000000, "Global write rate limit (messages/sec)")
	flag.IntVar(&Config.batchSize, "message-batch-size", 1, "Messages per batch")
	flag.StringVar(&Config.compressionName, "compression", "none", "Message compression: none, gzip, snappy")
	flag.StringVar(&Config.requiredAcksName, "required-acks", "local", "RequiredAcks config: none, local, all")
	flag.BoolVar(&Config.noop, "noop", false, "Test message generation performance (does not connect to Kafka)")
	flag.IntVar(&Config.workers, "workers", 1, "Number of workers")
	flag.IntVar(&Config.writersPerWorker, "writers-per-worker", 5, "Number of writer (Kafka producer) goroutines per worker")
	brokerString := flag.String("brokers", "localhost:9092", "Comma delimited list of Kafka brokers")
	flag.Parse()

	Config.brokers = strings.Split(*brokerString, ",")

	switch Config.compressionName {
	case "gzip":
		Config.compression = sarama.CompressionGZIP
	case "snappy":
		Config.compression = sarama.CompressionSnappy
	case "none":
		Config.compression = sarama.CompressionNone
	default:
		fmt.Printf("Invalid compression option: %s\n", Config.compressionName)
		os.Exit(1)
	}

	switch Config.requiredAcksName {
	case "none":
		Config.requiredAcks = sarama.NoResponse
	case "local":
		Config.requiredAcks = sarama.WaitForLocal
	case "all":
		Config.requiredAcks = sarama.WaitForAll
	default:
		fmt.Printf("Invalid required-acks option: %s\n", Config.requiredAcksName)
		os.Exit(1)
	}
}

func main() {
	if graphiteIp != "" {
		go graphiteWriter()
	}

	// Print Sangrenel startup info.
	fmt.Printf("\nStarting %d client workers, %d writers per worker\n", Config.workers, Config.writersPerWorker)
	fmt.Printf("Message size %d bytes, %d message limit per batch\n", Config.msgSize, Config.batchSize)
	fmt.Printf("Compression: %s, RequiredAcks: %s\n",
		Config.compressionName, Config.requiredAcksName)

	t := tachymeter.New(&tachymeter.Config{Size: 300000, Safe: true})

	// Start client workers.
	for i := 0; i < Config.workers; i++ {
		go worker(i+1, t)
	}

	var currSentCnt, lastSentCnt uint64
	var currErrCnt, lastErrCnt uint64

	interval := 5 * time.Second
	ticker := time.Tick(interval)
	start := time.Now()

	for {
		<-ticker

		intervalTime := time.Since(start).Seconds()

		// Set tachymeter wall time.
		t.SetWallTime(time.Since(start))

		// Get the sent count from the last interval, then the delta
		// (sentSinceLastInterval) between the current and last interval.
		lastSentCnt = currSentCnt
		currSentCnt = atomic.LoadUint64(&sentCnt)
		sentSinceLastInterval := currSentCnt - lastSentCnt

		outputBytes, outputString := calcOutput(intervalTime, sentSinceLastInterval)

		// Update error counters.
		lastErrCnt = currErrCnt
		currErrCnt = atomic.LoadUint64(&errCnt)
		errSinceLastInterval := currErrCnt - lastErrCnt

		errRate := (float64(errSinceLastInterval) / float64(sentSinceLastInterval)) * 100

		// Summarize tachymeter data.
		stats := t.Calc()

		// Update the metrics map for the Graphite writer.
		metrics["rate"] = float64(sentSinceLastInterval) / intervalTime
		metrics["error_rate"] = errRate
		metrics["output"] = outputBytes
		metrics["p99"] = (float64(stats.Time.P99.Nanoseconds()) / 1000) / 1000
		metrics["timestamp"] = float64(time.Now().Unix())
		// Ship metrics if configured.
		if graphiteIp != "" {
			metricsOutgoing <- metrics
		}

		// Write output stats.
		fmt.Println()
		log.Printf("[ topic: %s ]\n", Config.topic)
		fmt.Printf("> Messages: %s @ %.0f msgs/sec. | error rate %.2f%%\n",
			outputString,
			metrics["rate"],
			metrics["error_rate"])

		if !Config.noop {
			fmt.Printf("> Batches: %.2f batches/sec. | %s p99 | %s HMean | %s Min | %s Max\n",
				stats.Rate.Second, round(stats.Time.P99), round(stats.Time.HMean), round(stats.Time.Min), round(stats.Time.Max))

			fmt.Println(stats.Histogram.String(25))

			// Check if the tacymeter size needs to be increased
			// to avoid sampling. Otherwise, just reset it.
			if int(sentSinceLastInterval) > len(t.Times) {
				newTachy := tachymeter.New(&tachymeter.Config{Size: int(2 * sentSinceLastInterval)})
				// This is actually dangerous;
				// this could swap in a tachy with unlocked
				// mutexes while the current one has locks held.
				*t = *newTachy
			} else {
				t.Reset()
			}
		}

		// Reset interval time.
		start = time.Now()
	}
}

// worker is a high level producer unit and holds a single
// Kafka client. The worker's Kafka client is shared by n (Config.writersPerWorker)
// writer instances that perform the message generation and writing.
func worker(n int, t *tachymeter.Tachymeter) {
	switch Config.noop {
	case false:
		cId := "worker_" + strconv.Itoa(n)

		conf := sarama.NewConfig()
		conf.Producer.Compression = Config.compression
		conf.Producer.Return.Successes = true
		conf.Producer.RequiredAcks = Config.requiredAcks
		conf.Producer.Flush.MaxMessages = Config.batchSize
		conf.Producer.MaxMessageBytes = Config.msgSize + 50

		client, err := sarama.NewClient(Config.brokers, conf)
		if err != nil {
			log.Println(err)
			os.Exit(1)
		} else {
			log.Printf("%s connected\n", cId)
		}

		for i := 0; i < Config.writersPerWorker; i++ {
			go writer(client, t)
		}
	// If noop, we're not creating connections at all.
	// Just generate messages and burn CPU.
	default:
		for i := 0; i < Config.writersPerWorker; i++ {
			go dummyWriter(t)
		}
	}

	wait := make(chan bool)
	<-wait
}

// writer generates random messages and write to Kafka.
// Each wrtier belongs to a parent worker. Writers
// throttle writes according to a global rate limiter
// and report write throughput statistics up through
// a shared tachymeter.
func writer(c sarama.Client, t *tachymeter.Tachymeter) {
	// Init the producer.
	producer, err := sarama.NewSyncProducerFromClient(c)
	if err != nil {
		log.Println(err.Error())
	}
	defer producer.Close()

	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msgBatch := make([]*sarama.ProducerMessage, 0, Config.batchSize)

	for {
		// Message rate limiting works by having all writer loops incrementing
		// a global counter and tracking the aggregate per-second progress.
		// If the configured rate is met, the worker will sleep
		// for the remainder of the 1 second window.
		intervalEnd := time.Now().Add(time.Second)
		countStart := atomic.LoadUint64(&sentCnt)

		var sendTime time.Time

		for {
			// Break if the global rate limit was met, or, if
			// we'd exceed it assuming all writers wrote a max batch size
			// for this interval.
			intervalSent := atomic.LoadUint64(&sentCnt) - countStart
			sendEstimate := intervalSent + uint64(Config.batchSize*Config.workers*(Config.writersPerWorker-1))
			if sendEstimate >= Config.msgRate {
				break
			}

			// Estimate the batch size. This should shrink
			// if we're near the rate limit. Estimated batch size =
			// amount left to send for this interval / number of writers
			// we have available to send this amount. If the estimate
			// is lower than the configured batch size, send that amount
			// instead.
			toSend := (Config.msgRate - intervalSent) / uint64((Config.workers * Config.writersPerWorker))
			n := int(math.Min(float64(toSend), float64(Config.batchSize)))

			for i := 0; i < n; i++ {
				// Gen message.
				msgData := make([]byte, Config.msgSize)
				randMsg(msgData, *generator)
				msg := &sarama.ProducerMessage{Topic: Config.topic, Value: sarama.ByteEncoder(msgData)}
				// Append to batch.
				msgBatch = append(msgBatch, msg)
			}

			sendTime = time.Now()
			err = producer.SendMessages(msgBatch)
			if err != nil {
				// Sarama returns a ProducerErrors, which is a slice
				// of errors per message errored. Use this count
				// to establish an error rate.
				atomic.AddUint64(&errCnt, uint64(len(err.(sarama.ProducerErrors))))
			}

			t.AddTime(time.Since(sendTime))
			atomic.AddUint64(&sentCnt, uint64(len(msgBatch)))

			msgBatch = msgBatch[:0]
		}

		// If the global per-second rate limit was met,
		// the inner loop breaks and the outer loop sleeps for the interval remainder.
		time.Sleep(intervalEnd.Sub(time.Now()))
	}
}

// dummyWriter is initialized by the worker(s) if Config.noop is True.
// dummyWriter performs the message generation step of the normal writer,
// but doesn't connect to / attempt to send anything to Kafka. This is used
// purely for testing message generation performance.
func dummyWriter(t *tachymeter.Tachymeter) {
	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msgBatch := make([]*sarama.ProducerMessage, 0, Config.batchSize)

	for {
		for i := 0; i < Config.batchSize; i++ {
			// Gen message.
			msgData := make([]byte, Config.msgSize)
			randMsg(msgData, *generator)
			msg := &sarama.ProducerMessage{Topic: Config.topic, Value: sarama.ByteEncoder(msgData)}
			// Append to batch.
			msgBatch = append(msgBatch, msg)
		}

		atomic.AddUint64(&sentCnt, uint64(len(msgBatch)))
		t.AddTime(time.Duration(0))

		msgBatch = msgBatch[:0]
	}
}

// randMsg returns a random message generated from the chars byte slice.
// Message length of m bytes as defined by Config.msgSize.
func randMsg(m []byte, generator rand.Rand) {
	for i := range m {
		m[i] = chars[generator.Intn(len(chars))]
	}
}

// calcOutput takes a duration t and messages sent
// and returns message rates in human readable network speeds.
func calcOutput(t float64, n uint64) (float64, string) {
	m := (float64(n) / t) * float64(Config.msgSize)
	var o string
	switch {
	case m >= 131072:
		o = strconv.FormatFloat(m/131072, 'f', 0, 64) + "Mb/sec"
	case m < 131072:
		o = strconv.FormatFloat(m/1024, 'f', 0, 64) + "KB/sec"
	}
	return m, o
}

func round(t time.Duration) time.Duration {
	return t / 1000 * 1000
}
<commit_msg>fixes rate limiter with concurrent workers<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/jamiealquiza/tachymeter"
	"gopkg.in/Shopify/sarama.v1"
)

type config struct {
	brokers          []string
	topic            string
	msgSize          int
	msgRate          uint64
	batchSize        int
	compression      sarama.CompressionCodec
	compressionName  string
	requiredAcks     sarama.RequiredAcks
	requiredAcksName string
	workers          int
	writersPerWorker int
	noop             bool
}

var (
	Config = &config{}

	// Character selection for random messages.
	chars = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$^&*(){}][:<>.")

	// Counters / misc.
	sentCnt uint64
	errCnt  uint64
)

func init() {
	flag.StringVar(&Config.topic, "topic", "sangrenel", "Kafka topic to produce to")
	flag.IntVar(&Config.msgSize, "message-size", 300, "Message size (bytes)")
	flag.Uint64Var(&Config.msgRate, "produce-rate", 100000000, "Global write rate limit (messages/sec)")
	flag.IntVar(&Config.batchSize, "message-batch-size", 1, "Messages per batch")
	flag.StringVar(&Config.compressionName, "compression", "none", "Message compression: none, gzip, snappy")
	flag.StringVar(&Config.requiredAcksName, "required-acks", "local", "RequiredAcks config: none, local, all")
	flag.BoolVar(&Config.noop, "noop", false, "Test message generation performance (does not connect to Kafka)")
	flag.IntVar(&Config.workers, "workers", 1, "Number of workers")
	flag.IntVar(&Config.writersPerWorker, "writers-per-worker", 5, "Number of writer (Kafka producer) goroutines per worker")
	brokerString := flag.String("brokers", "localhost:9092", "Comma delimited list of Kafka brokers")
	flag.Parse()

	Config.brokers = strings.Split(*brokerString, ",")

	switch Config.compressionName {
	case "gzip":
		Config.compression = sarama.CompressionGZIP
	case "snappy":
		Config.compression = sarama.CompressionSnappy
	case "none":
		Config.compression = sarama.CompressionNone
	default:
		fmt.Printf("Invalid compression option: %s\n", Config.compressionName)
		os.Exit(1)
	}

	switch Config.requiredAcksName {
	case "none":
		Config.requiredAcks = sarama.NoResponse
	case "local":
		Config.requiredAcks = sarama.WaitForLocal
	case "all":
		Config.requiredAcks = sarama.WaitForAll
	default:
		fmt.Printf("Invalid required-acks option: %s\n", Config.requiredAcksName)
		os.Exit(1)
	}
}

func main() {
	if graphiteIp != "" {
		go graphiteWriter()
	}

	// Print Sangrenel startup info.
	fmt.Printf("\nStarting %d client workers, %d writers per worker\n", Config.workers, Config.writersPerWorker)
	fmt.Printf("Message size %d bytes, %d message limit per batch\n", Config.msgSize, Config.batchSize)
	fmt.Printf("Compression: %s, RequiredAcks: %s\n",
		Config.compressionName, Config.requiredAcksName)

	t := tachymeter.New(&tachymeter.Config{Size: 300000, Safe: true})

	// Start client workers.
	for i := 0; i < Config.workers; i++ {
		go worker(i+1, t)
	}

	var currSentCnt, lastSentCnt uint64
	var currErrCnt, lastErrCnt uint64

	interval := 5 * time.Second
	ticker := time.Tick(interval)
	start := time.Now()

	for {
		<-ticker

		intervalTime := time.Since(start).Seconds()

		// Set tachymeter wall time.
		t.SetWallTime(time.Since(start))

		// Get the sent count from the last interval, then the delta
		// (sentSinceLastInterval) between the current and last interval.
		lastSentCnt = currSentCnt
		currSentCnt = atomic.LoadUint64(&sentCnt)
		sentSinceLastInterval := currSentCnt - lastSentCnt

		outputBytes, outputString := calcOutput(intervalTime, sentSinceLastInterval)

		// Update error counters.
		lastErrCnt = currErrCnt
		currErrCnt = atomic.LoadUint64(&errCnt)
		errSinceLastInterval := currErrCnt - lastErrCnt

		errRate := (float64(errSinceLastInterval) / float64(sentSinceLastInterval)) * 100

		// Summarize tachymeter data.
		stats := t.Calc()

		// Update the metrics map for the Graphite writer.
		metrics["rate"] = float64(sentSinceLastInterval) / intervalTime
		metrics["error_rate"] = errRate
		metrics["output"] = outputBytes
		metrics["p99"] = (float64(stats.Time.P99.Nanoseconds()) / 1000) / 1000
		metrics["timestamp"] = float64(time.Now().Unix())
		// Ship metrics if configured.
		if graphiteIp != "" {
			metricsOutgoing <- metrics
		}

		// Write output stats.
		fmt.Println()
		log.Printf("[ topic: %s ]\n", Config.topic)
		fmt.Printf("> Messages: %s @ %.0f msgs/sec. | error rate %.2f%%\n",
			outputString,
			metrics["rate"],
			metrics["error_rate"])

		if !Config.noop {
			fmt.Printf("> Batches: %.2f batches/sec. | %s p99 | %s HMean | %s Min | %s Max\n",
				stats.Rate.Second, round(stats.Time.P99), round(stats.Time.HMean), round(stats.Time.Min), round(stats.Time.Max))

			fmt.Println(stats.Histogram.String(25))

			// Check if the tacymeter size needs to be increased
			// to avoid sampling. Otherwise, just reset it.
			if int(sentSinceLastInterval) > len(t.Times) {
				newTachy := tachymeter.New(&tachymeter.Config{Size: int(2 * sentSinceLastInterval)})
				// This is actually dangerous;
				// this could swap in a tachy with unlocked
				// mutexes while the current one has locks held.
				*t = *newTachy
			} else {
				t.Reset()
			}
		}

		// Reset interval time.
		start = time.Now()
	}
}

// worker is a high level producer unit and holds a single
// Kafka client. The worker's Kafka client is shared by n (Config.writersPerWorker)
// writer instances that perform the message generation and writing.
func worker(n int, t *tachymeter.Tachymeter) {
	switch Config.noop {
	case false:
		cId := "worker_" + strconv.Itoa(n)

		conf := sarama.NewConfig()
		conf.Producer.Compression = Config.compression
		conf.Producer.Return.Successes = true
		conf.Producer.RequiredAcks = Config.requiredAcks
		conf.Producer.Flush.MaxMessages = Config.batchSize
		conf.Producer.MaxMessageBytes = Config.msgSize + 50

		client, err := sarama.NewClient(Config.brokers, conf)
		if err != nil {
			log.Println(err)
			os.Exit(1)
		} else {
			log.Printf("%s connected\n", cId)
		}

		for i := 0; i < Config.writersPerWorker; i++ {
			go writer(client, t)
		}
	// If noop, we're not creating connections at all.
	// Just generate messages and burn CPU.
	default:
		for i := 0; i < Config.writersPerWorker; i++ {
			go dummyWriter(t)
		}
	}

	wait := make(chan bool)
	<-wait
}

// writer generates random messages and write to Kafka.
// Each wrtier belongs to a parent worker. Writers
// throttle writes according to a global rate limiter
// and report write throughput statistics up through
// a shared tachymeter.
func writer(c sarama.Client, t *tachymeter.Tachymeter) {
	// Init the producer.
	producer, err := sarama.NewSyncProducerFromClient(c)
	if err != nil {
		log.Println(err.Error())
	}
	defer producer.Close()

	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msgBatch := make([]*sarama.ProducerMessage, 0, Config.batchSize)

	for {
		// Message rate limiting works by having all writer loops incrementing
		// a global counter and tracking the aggregate per-second progress.
		// If the configured rate is met, the worker will sleep
		// for the remainder of the 1 second window.
		intervalEnd := time.Now().Add(time.Second)
		countStart := atomic.LoadUint64(&sentCnt)

		var sendTime time.Time
		var intervalSent uint64

		for {
			// Estimate the batch size. This should shrink
			// if we're near the rate limit. Estimated batch size =
			// amount left to send for this interval / number of writers
			// we have available to send this amount. If the estimate
			// is lower than the configured batch size, send that amount
			// instead.
			toSend := (Config.msgRate - intervalSent) / uint64((Config.workers * Config.writersPerWorker))
			n := int(math.Min(float64(toSend), float64(Config.batchSize)))

			for i := 0; i < n; i++ {
				// Gen message.
				msgData := make([]byte, Config.msgSize)
				randMsg(msgData, *generator)
				msg := &sarama.ProducerMessage{Topic: Config.topic, Value: sarama.ByteEncoder(msgData)}
				// Append to batch.
				msgBatch = append(msgBatch, msg)
			}

			sendTime = time.Now()
			err = producer.SendMessages(msgBatch)
			if err != nil {
				// Sarama returns a ProducerErrors, which is a slice
				// of errors per message errored. Use this count
				// to establish an error rate.
				atomic.AddUint64(&errCnt, uint64(len(err.(sarama.ProducerErrors))))
			}

			t.AddTime(time.Since(sendTime))
			atomic.AddUint64(&sentCnt, uint64(len(msgBatch)))

			msgBatch = msgBatch[:0]

			intervalSent = atomic.LoadUint64(&sentCnt) - countStart

			// Break if the global rate limit was met, or, if
			// we'd exceed it assuming all writers wrote a max batch size
			// for this interval. This is conditionally ran at the beginning
			// or the end of the interval depending on worker concurrency.
			sendEstimate := intervalSent + uint64((Config.batchSize*Config.workers*Config.writersPerWorker)-Config.writersPerWorker)
			if sendEstimate >= Config.msgRate {
				break
			}
		}

		// If the global per-second rate limit was met,
		// the inner loop breaks and the outer loop sleeps for the interval remainder.
		time.Sleep(intervalEnd.Sub(time.Now()))
	}
}

// dummyWriter is initialized by the worker(s) if Config.noop is True.
// dummyWriter performs the message generation step of the normal writer,
// but doesn't connect to / attempt to send anything to Kafka. This is used
// purely for testing message generation performance.
func dummyWriter(t *tachymeter.Tachymeter) {
	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msgBatch := make([]*sarama.ProducerMessage, 0, Config.batchSize)

	for {
		for i := 0; i < Config.batchSize; i++ {
			// Gen message.
			msgData := make([]byte, Config.msgSize)
			randMsg(msgData, *generator)
			msg := &sarama.ProducerMessage{Topic: Config.topic, Value: sarama.ByteEncoder(msgData)}
			// Append to batch.
			msgBatch = append(msgBatch, msg)
		}

		atomic.AddUint64(&sentCnt, uint64(len(msgBatch)))
		t.AddTime(time.Duration(0))

		msgBatch = msgBatch[:0]
	}
}

// randMsg returns a random message generated from the chars byte slice.
// Message length of m bytes as defined by Config.msgSize.
func randMsg(m []byte, generator rand.Rand) {
	for i := range m {
		m[i] = chars[generator.Intn(len(chars))]
	}
}

// calcOutput takes a duration t and messages sent
// and returns message rates in human readable network speeds.
func calcOutput(t float64, n uint64) (float64, string) {
	m := (float64(n) / t) * float64(Config.msgSize)
	var o string
	switch {
	case m >= 131072:
		o = strconv.FormatFloat(m/131072, 'f', 0, 64) + "Mb/sec"
	case m < 131072:
		o = strconv.FormatFloat(m/1024, 'f', 0, 64) + "KB/sec"
	}
	return m, o
}

func round(t time.Duration) time.Duration {
	return t / 1000 * 1000
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"encoding/json"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"syscall"
	"strconv"
	"time"
	"io"
	"net/http"
	"github.com/hoisie/redis"
)

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
	MAX_UDP_PACKET_SIZE     = 512
)

var (
	serviceAddress   = flag.String("address", ":8125", "UDP service address")
	expiryInterval   = flag.Int64("expiry-interval", 1, "Expiry interval (seconds)")
	debug            = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion      = flag.Bool("version", false, "print version string")
)

const (
	ACTION_SET_WARN = "set-warn"
	ACTION_SET_ERR = "set-err"
	ACTION_BEAT = "beat"
)

const (
	STATE_PAUSED  = "paused"
	STATE_OK      = "ok"
	STATE_WARNING = "warning"
	STATE_ERROR   = "error"
)

type Cmd struct {
	Action   string
	Service  string
	Value    int
}

type Service struct {
	Name           string
	LastValue      int
	LastBeat       int64
	LastUpdated    int64
	WarningTimeout int64
	ErrorTimeout   int64
	State          string
}

var (
	In       = make(chan *Cmd, MAX_UNPROCESSED_PACKETS)
	signalchan chan os.Signal
	client redis.Client
)

func now() int64 { return time.Now().Unix() }

func getOrCreate(name string) (*Service, *Service) {
	service := &Service{
		Name: name,
		LastValue: -1,
		LastBeat: -1,
		LastUpdated: -1,
		WarningTimeout: -1,
		ErrorTimeout: -1,
		State: STATE_PAUSED,
	}

	if data, err := client.Get("lb.service." + name); err == nil {
		json.Unmarshal(data, &service)
	}
	var ref = *service
	return service, &ref
}

func (s *Service)GetExpiry(timeout int64) int64 {
	if timeout <= 0 {
		return 0
	}
	return s.LastBeat + timeout
}

func (s *Service)GetNextExpiry(ts int64) int64 {
	var next int64 = 0
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && warningExpiry > ts && (next == 0 || warningExpiry < next) {
		next = warningExpiry
	}
	if errorExpiry > 0 && errorExpiry > ts && (next == 0 || errorExpiry < next) {
		next = errorExpiry
	}
	return next
}

func (s *Service) UpdateState(ts int64) {
	s.State = STATE_OK
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && ts >= warningExpiry {
		s.State = STATE_WARNING
	}
	if errorExpiry > 0 && ts >= errorExpiry {
		s.State = STATE_ERROR
	}
}

func updateExpiry(service *Service, ts int64) {
	if service.State != STATE_PAUSED {
		if expiry := service.GetNextExpiry(ts); expiry > 0 {
			client.Zadd("lb.expiry", []byte(service.Name), float64(expiry))
			return
		}
	}
	client.Zrem("lb.expiry", []byte(service.Name))
}

func (s *Service)Save(ref *Service, ts int64) {
	if *s != *ref {
		if s.State != ref.State {
			log.Printf("service %s, state %s -> %s", s.Name, ref.State, s.State)
		}
		if s.WarningTimeout != ref.WarningTimeout {
			log.Printf("service %s, warn %d -> %d", s.Name, ref.WarningTimeout, s.WarningTimeout)
		}
		if s.ErrorTimeout != ref.ErrorTimeout {
			log.Printf("service %s, err %d -> %d", s.Name, ref.ErrorTimeout, s.ErrorTimeout)
		}
		s.LastUpdated = ts
		b, _ := json.Marshal(s)
		client.Set("lb.service." + s.Name, b)
		if ref.LastUpdated < 0 {
			client.Sadd("lb.services.all", []byte(s.Name))
		}
	}
}

func monitor() {
	period := time.Duration(*expiryInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			return
		case <-ticker.C:
			var ts = now()
			if expired, err := client.Zrangebyscore("lb.expiry", 0, float64(now())); err == nil {
				for _, elem := range expired {
					var service, ref = getOrCreate(string(elem))
					service.UpdateState(ts)
					service.Save(ref, ts)
					updateExpiry(service, ts)
				}
			}
		case s := <-In:
			var ts = now()
			var service, ref = getOrCreate(s.Service)
			switch s.Action {
			case ACTION_SET_WARN:
				service.WarningTimeout = int64(s.Value)
			case ACTION_SET_ERR:
				service.ErrorTimeout = int64(s.Value)
			case ACTION_BEAT:
				service.LastBeat = ts
			}
			service.UpdateState(ts)
			service.Save(ref, ts)
			updateExpiry(service, ts)
		}
	}
}

var packetRegexp = regexp.MustCompile("^([^:]+)\\.(beat|warn|err):(-?[0-9]+)\\|(g|c|ms)(\\|@([0-9\\.]+))?\n?$")

func parseMessage(data []byte) []*Cmd {
	var output []*Cmd
	for _, line := range bytes.Split(data, []byte("\n")) {
		if len(line) == 0 {
			continue
		}

		item := packetRegexp.FindSubmatch(line)
		if len(item) == 0 {
			continue
		}

		var value int
		modifier := string(item[4])
		switch modifier {
		case "c":
			var vali, err = strconv.ParseInt(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseInt %s - %s", item[3], err)
				continue
			}
			value = int(vali)
		default:
			var valu, err = strconv.ParseUint(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseUint %s - %s", item[3], err)
				continue
			}
			value = int(valu)
		}
		var action string
		switch string(item[2]) {
		case "warn":
			action = ACTION_SET_WARN
		case "err":
			action = ACTION_SET_ERR
		case "beat":
			action = ACTION_BEAT
		}

		packet := &Cmd{
			Action: action,
			Service: string(item[1]),
			Value:    value,
		}
		output = append(output, packet)
	}
	return output
}

func udpListener() {
	address, _ := net.ResolveUDPAddr("udp", *serviceAddress)
	log.Printf("listening on %s", address)
	listener, err := net.ListenUDP("udp", address)
	if err != nil {
		log.Fatalf("ERROR: ListenUDP - %s", err)
	}
	defer listener.Close()

	message := make([]byte, MAX_UDP_PACKET_SIZE)
	for {
		n, remaddr, err := listener.ReadFromUDP(message)
		if err != nil {
			log.Printf("ERROR: reading UDP packet from %+v - %s", remaddr, err)
			continue
		}

		for _, p := range parseMessage(message[:n]) {
			In <- p
		}
	}
}

func DashboardHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		buffer.WriteString(fmt.Sprintf("%-40s%s\n", service.Name, service.State))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("\nwarnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func StatusHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("warnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func httpServer() {
	http.Handle("/", http.HandlerFunc(DashboardHandler))
	http.Handle("/status", http.HandlerFunc(StatusHandler))
        http.ListenAndServe(":8080", nil)
}

func main() {
	flag.Parse()
	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	signalchan = make(chan os.Signal, 1)
	signal.Notify(signalchan, syscall.SIGTERM)

	go httpServer()
	go udpListener()
	monitor()
}
<commit_msg>Added trigger handler<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"encoding/json"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"syscall"
	"strconv"
	"time"
	"io"
	"net/http"
	"github.com/hoisie/redis"
)

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
	MAX_UDP_PACKET_SIZE     = 512
)

var (
	serviceAddress   = flag.String("address", ":8125", "UDP service address")
	expiryInterval   = flag.Int64("expiry-interval", 1, "Expiry interval (seconds)")
	debug            = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion      = flag.Bool("version", false, "print version string")
)

const (
	ACTION_SET_WARN = "set-warn"
	ACTION_SET_ERR = "set-err"
	ACTION_BEAT = "beat"
)

const (
	STATE_PAUSED  = "paused"
	STATE_OK      = "ok"
	STATE_WARNING = "warning"
	STATE_ERROR   = "error"
)

type Cmd struct {
	Action   string
	Service  string
	Value    int
}

type Service struct {
	Name           string
	LastValue      int
	LastBeat       int64
	LastUpdated    int64
	WarningTimeout int64
	ErrorTimeout   int64
	State          string
}

var (
	In       = make(chan *Cmd, MAX_UNPROCESSED_PACKETS)
	signalchan chan os.Signal
	client redis.Client
)

func now() int64 { return time.Now().Unix() }

func getOrCreate(name string) (*Service, *Service) {
	service := &Service{
		Name: name,
		LastValue: -1,
		LastBeat: -1,
		LastUpdated: -1,
		WarningTimeout: -1,
		ErrorTimeout: -1,
		State: STATE_PAUSED,
	}

	if data, err := client.Get("lb.service." + name); err == nil {
		json.Unmarshal(data, &service)
	}
	var ref = *service
	return service, &ref
}

func (s *Service)GetExpiry(timeout int64) int64 {
	if timeout <= 0 {
		return 0
	}
	return s.LastBeat + timeout
}

func (s *Service)GetNextExpiry(ts int64) int64 {
	var next int64 = 0
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && warningExpiry > ts && (next == 0 || warningExpiry < next) {
		next = warningExpiry
	}
	if errorExpiry > 0 && errorExpiry > ts && (next == 0 || errorExpiry < next) {
		next = errorExpiry
	}
	return next
}

func (s *Service) UpdateState(ts int64) {
	s.State = STATE_OK
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && ts >= warningExpiry {
		s.State = STATE_WARNING
	}
	if errorExpiry > 0 && ts >= errorExpiry {
		s.State = STATE_ERROR
	}
}

func updateExpiry(service *Service, ts int64) {
	if service.State != STATE_PAUSED {
		if expiry := service.GetNextExpiry(ts); expiry > 0 {
			client.Zadd("lb.expiry", []byte(service.Name), float64(expiry))
			return
		}
	}
	client.Zrem("lb.expiry", []byte(service.Name))
}

func (s *Service)Save(ref *Service, ts int64) {
	if *s != *ref {
		if s.State != ref.State {
			log.Printf("service %s, state %s -> %s", s.Name, ref.State, s.State)
		}
		if s.WarningTimeout != ref.WarningTimeout {
			log.Printf("service %s, warn %d -> %d", s.Name, ref.WarningTimeout, s.WarningTimeout)
		}
		if s.ErrorTimeout != ref.ErrorTimeout {
			log.Printf("service %s, err %d -> %d", s.Name, ref.ErrorTimeout, s.ErrorTimeout)
		}
		s.LastUpdated = ts
		b, _ := json.Marshal(s)
		client.Set("lb.service." + s.Name, b)
		if ref.LastUpdated < 0 {
			client.Sadd("lb.services.all", []byte(s.Name))
		}
	}
}

func monitor() {
	period := time.Duration(*expiryInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			return
		case <-ticker.C:
			var ts = now()
			if expired, err := client.Zrangebyscore("lb.expiry", 0, float64(now())); err == nil {
				for _, elem := range expired {
					var service, ref = getOrCreate(string(elem))
					service.UpdateState(ts)
					service.Save(ref, ts)
					updateExpiry(service, ts)
				}
			}
		case s := <-In:
			var ts = now()
			var service, ref = getOrCreate(s.Service)
			switch s.Action {
			case ACTION_SET_WARN:
				service.WarningTimeout = int64(s.Value)
			case ACTION_SET_ERR:
				service.ErrorTimeout = int64(s.Value)
			case ACTION_BEAT:
				service.LastBeat = ts
			}
			service.UpdateState(ts)
			service.Save(ref, ts)
			updateExpiry(service, ts)
		}
	}
}

var packetRegexp = regexp.MustCompile("^([^:]+)\\.(beat|warn|err):(-?[0-9]+)\\|(g|c|ms)(\\|@([0-9\\.]+))?\n?$")

func parseMessage(data []byte) []*Cmd {
	var output []*Cmd
	for _, line := range bytes.Split(data, []byte("\n")) {
		if len(line) == 0 {
			continue
		}

		item := packetRegexp.FindSubmatch(line)
		if len(item) == 0 {
			continue
		}

		var value int
		modifier := string(item[4])
		switch modifier {
		case "c":
			var vali, err = strconv.ParseInt(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseInt %s - %s", item[3], err)
				continue
			}
			value = int(vali)
		default:
			var valu, err = strconv.ParseUint(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseUint %s - %s", item[3], err)
				continue
			}
			value = int(valu)
		}
		var action string
		switch string(item[2]) {
		case "warn":
			action = ACTION_SET_WARN
		case "err":
			action = ACTION_SET_ERR
		case "beat":
			action = ACTION_BEAT
		}

		packet := &Cmd{
			Action: action,
			Service: string(item[1]),
			Value:    value,
		}
		output = append(output, packet)
	}
	return output
}

func udpListener() {
	address, _ := net.ResolveUDPAddr("udp", *serviceAddress)
	log.Printf("listening on %s", address)
	listener, err := net.ListenUDP("udp", address)
	if err != nil {
		log.Fatalf("ERROR: ListenUDP - %s", err)
	}
	defer listener.Close()

	message := make([]byte, MAX_UDP_PACKET_SIZE)
	for {
		n, remaddr, err := listener.ReadFromUDP(message)
		if err != nil {
			log.Printf("ERROR: reading UDP packet from %+v - %s", remaddr, err)
			continue
		}

		for _, p := range parseMessage(message[:n]) {
			In <- p
		}
	}
}

func DashboardHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		buffer.WriteString(fmt.Sprintf("%-40s%s\n", service.Name, service.State))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("\nwarnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func StatusHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("warnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func TriggerHandler(c http.ResponseWriter, r *http.Request) {
	var err = r.ParseForm()
	if err != nil {
		log.Print("error parsing form ", err)
		return
	}

	var errtmo, warntmo = r.FormValue("err-tmo"), r.FormValue("warn-tmo")

	In <- &Cmd{
		Action:  ACTION_BEAT,
		Service: string("path"),
		Value:   1,
	}

	if val, err := strconv.Atoi(errtmo); err == nil {
		In <- &Cmd{
			Action:  ACTION_SET_ERR,
			Service: string("path"),
			Value:   val,
		}
	}

	if val, err := strconv.Atoi(warntmo); err == nil {
		In <- &Cmd{
			Action:  ACTION_SET_WARN,
			Service: string("path"),
			Value:   val,
		}
	}

        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", "3")
        io.WriteString(c, "ok\n")
}

func httpServer() {
	http.Handle("/", http.HandlerFunc(DashboardHandler))
	http.Handle("/status", http.HandlerFunc(StatusHandler))
	http.Handle("/trigger", http.HandlerFunc(TriggerHandler))
        http.ListenAndServe(":8080", nil)
}

func main() {
	flag.Parse()
	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	signalchan = make(chan os.Signal, 1)
	signal.Notify(signalchan, syscall.SIGTERM)

	go httpServer()
	go udpListener()
	monitor()
}
<|endoftext|>


<commit_before>package lua

/*
#cgo CFLAGS: -Ilua

#include <lua.h>
#include <lualib.h>
#include <stdlib.h>
*/
import "C"

import (
	"reflect"
	"unsafe"
)

// Type of allocation functions to use with NewStateAlloc
type Alloc func(ptr unsafe.Pointer, osize uint, nsize uint) unsafe.Pointer

// This is the type of go function that can be registered as lua functions
type LuaGoFunction func(L *State) int

// Wrapper to keep cgo from complaining about incomplete ptr type
//export State
type State struct {
	// Wrapped lua_State object
	s *C.lua_State

	// Registry of go object that have been pushed to Lua VM
	registry []interface{}

	// Freelist for funcs indices, to allow for freeing
	freeIndices []uint
}

//export golua_callgofunction
func golua_callgofunction(L interface{}, fid uint) int {
	L1 := L.(*State)
	if fid < 0 {
		panic(&LuaError{0, "Requested execution of an unknown function", L1.StackTrace()})
	}
	f := L1.registry[fid].(LuaGoFunction)
	return f(L1)
}

//export golua_interface_newindex_callback
func golua_interface_newindex_callback(Li interface{}, iid uint, field_name_cstr *C.char) int {
	L := Li.(*State)
	iface := L.registry[iid]
	ifacevalue := reflect.ValueOf(iface).Elem()

	field_name := C.GoString(field_name_cstr)

	fval := ifacevalue.FieldByName(field_name)

	if fval.Kind() == reflect.Ptr {
		fval = fval.Elem()
	}

	luatype := LuaValType(C.lua_type(L.s, 3))

	switch fval.Kind() {
	case reflect.Bool:
		if luatype == LUA_TBOOLEAN {
			fval.SetBool(int(C.lua_toboolean(L.s, 3)) != 0)
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.Int:
		fallthrough
	case reflect.Int8:
		fallthrough
	case reflect.Int16:
		fallthrough
	case reflect.Int32:
		fallthrough
	case reflect.Int64:
		if luatype == LUA_TNUMBER {
			fval.SetInt(int64(C.lua_tointeger(L.s, 3)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.Uint:
		fallthrough
	case reflect.Uint8:
		fallthrough
	case reflect.Uint16:
		fallthrough
	case reflect.Uint32:
		fallthrough
	case reflect.Uint64:
		if luatype == LUA_TNUMBER {
			fval.SetUint(uint64(C.lua_tointeger(L.s, 3)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.String:
		if luatype == LUA_TSTRING {
			fval.SetString(C.GoString(C.lua_tolstring(L.s, 3, nil)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.Float32:
		fallthrough
	case reflect.Float64:
		if luatype == LUA_TNUMBER {
			fval.SetFloat(float64(C.lua_tonumber(L.s, 3)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}
	}

	L.PushString("Unsupported type of field "+field_name+": "+fval.Type().String())
	return -1
}

//export golua_interface_index_callback
func golua_interface_index_callback(Li interface{}, iid uint, field_name *C.char) int {
	L := Li.(*State)
	iface := L.registry[iid]
	ifacevalue := reflect.ValueOf(iface).Elem()

	fval := ifacevalue.FieldByName(C.GoString(field_name))

	if fval.Kind() == reflect.Ptr {
		fval = fval.Elem()
	}

	switch fval.Kind() {
	case reflect.Bool:
		L.PushBoolean(fval.Bool())
		return 1

	case reflect.Int:
		fallthrough
	case reflect.Int8:
		fallthrough
	case reflect.Int16:
		fallthrough
	case reflect.Int32:
		fallthrough
	case reflect.Int64:
		L.PushInteger(fval.Int())
		return 1

	case reflect.Uint:
		fallthrough
	case reflect.Uint8:
		fallthrough
	case reflect.Uint16:
		fallthrough
	case reflect.Uint32:
		fallthrough
	case reflect.Uint64:
		L.PushInteger(int64(fval.Uint()))
		return 1

	case reflect.String:
		L.PushString(fval.String())
		return 1

	case reflect.Float32:
		fallthrough
	case reflect.Float64:
		L.PushNumber(fval.Float())
		return 1
	}

	L.PushString("Unsupported type of field: "+fval.Type().String())
	return -1
}

//export golua_gchook
func golua_gchook(L interface{}, id uint) int {
	L1 := L.(*State)
	L1.unregister(id)
	return 0
}

//export golua_callpanicfunction
func golua_callpanicfunction(L interface{}, id uint) int {
	L1 := L.(*State)
	f := L1.registry[id].(LuaGoFunction)
	return f(L1)
}

//export golua_idtointerface
func golua_idtointerface(id uint) interface{} {
	return id
}

//export golua_cfunctiontointerface
func golua_cfunctiontointerface(f *uintptr) interface{} {
	return f
}

//export golua_callallocf
func golua_callallocf(fp uintptr, ptr uintptr, osize uint, nsize uint) uintptr {
	return uintptr((*((*Alloc)(unsafe.Pointer(fp))))(unsafe.Pointer(ptr), osize, nsize))
}

//export go_panic_msghandler
func go_panic_msghandler(Li interface{}, z *C.char) {
	L := Li.(*State)
	s := C.GoString(z)

	panic(&LuaError{LUA_ERRERR, s, L.StackTrace()})
}
<commit_msg>Update golua.go<commit_after>package lua

/*
#cgo pkg-config: luajit

#include <lua.h>
#include <lualib.h>
#include <stdlib.h>
*/
import "C"

import (
	"reflect"
	"unsafe"
)

// Type of allocation functions to use with NewStateAlloc
type Alloc func(ptr unsafe.Pointer, osize uint, nsize uint) unsafe.Pointer

// This is the type of go function that can be registered as lua functions
type LuaGoFunction func(L *State) int

// Wrapper to keep cgo from complaining about incomplete ptr type
//export State
type State struct {
	// Wrapped lua_State object
	s *C.lua_State

	// Registry of go object that have been pushed to Lua VM
	registry []interface{}

	// Freelist for funcs indices, to allow for freeing
	freeIndices []uint
}

//export golua_callgofunction
func golua_callgofunction(L interface{}, fid uint) int {
	L1 := L.(*State)
	if fid < 0 {
		panic(&LuaError{0, "Requested execution of an unknown function", L1.StackTrace()})
	}
	f := L1.registry[fid].(LuaGoFunction)
	return f(L1)
}

//export golua_interface_newindex_callback
func golua_interface_newindex_callback(Li interface{}, iid uint, field_name_cstr *C.char) int {
	L := Li.(*State)
	iface := L.registry[iid]
	ifacevalue := reflect.ValueOf(iface).Elem()

	field_name := C.GoString(field_name_cstr)

	fval := ifacevalue.FieldByName(field_name)

	if fval.Kind() == reflect.Ptr {
		fval = fval.Elem()
	}

	luatype := LuaValType(C.lua_type(L.s, 3))

	switch fval.Kind() {
	case reflect.Bool:
		if luatype == LUA_TBOOLEAN {
			fval.SetBool(int(C.lua_toboolean(L.s, 3)) != 0)
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.Int:
		fallthrough
	case reflect.Int8:
		fallthrough
	case reflect.Int16:
		fallthrough
	case reflect.Int32:
		fallthrough
	case reflect.Int64:
		if luatype == LUA_TNUMBER {
			fval.SetInt(int64(C.lua_tointeger(L.s, 3)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.Uint:
		fallthrough
	case reflect.Uint8:
		fallthrough
	case reflect.Uint16:
		fallthrough
	case reflect.Uint32:
		fallthrough
	case reflect.Uint64:
		if luatype == LUA_TNUMBER {
			fval.SetUint(uint64(C.lua_tointeger(L.s, 3)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.String:
		if luatype == LUA_TSTRING {
			fval.SetString(C.GoString(C.lua_tolstring(L.s, 3, nil)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}

	case reflect.Float32:
		fallthrough
	case reflect.Float64:
		if luatype == LUA_TNUMBER {
			fval.SetFloat(float64(C.lua_tonumber(L.s, 3)))
			return 1
		} else {
			L.PushString("Wrong assignment to field "+field_name)
			return -1
		}
	}

	L.PushString("Unsupported type of field "+field_name+": "+fval.Type().String())
	return -1
}

//export golua_interface_index_callback
func golua_interface_index_callback(Li interface{}, iid uint, field_name *C.char) int {
	L := Li.(*State)
	iface := L.registry[iid]
	ifacevalue := reflect.ValueOf(iface).Elem()

	fval := ifacevalue.FieldByName(C.GoString(field_name))

	if fval.Kind() == reflect.Ptr {
		fval = fval.Elem()
	}

	switch fval.Kind() {
	case reflect.Bool:
		L.PushBoolean(fval.Bool())
		return 1

	case reflect.Int:
		fallthrough
	case reflect.Int8:
		fallthrough
	case reflect.Int16:
		fallthrough
	case reflect.Int32:
		fallthrough
	case reflect.Int64:
		L.PushInteger(fval.Int())
		return 1

	case reflect.Uint:
		fallthrough
	case reflect.Uint8:
		fallthrough
	case reflect.Uint16:
		fallthrough
	case reflect.Uint32:
		fallthrough
	case reflect.Uint64:
		L.PushInteger(int64(fval.Uint()))
		return 1

	case reflect.String:
		L.PushString(fval.String())
		return 1

	case reflect.Float32:
		fallthrough
	case reflect.Float64:
		L.PushNumber(fval.Float())
		return 1
	}

	L.PushString("Unsupported type of field: "+fval.Type().String())
	return -1
}

//export golua_gchook
func golua_gchook(L interface{}, id uint) int {
	L1 := L.(*State)
	L1.unregister(id)
	return 0
}

//export golua_callpanicfunction
func golua_callpanicfunction(L interface{}, id uint) int {
	L1 := L.(*State)
	f := L1.registry[id].(LuaGoFunction)
	return f(L1)
}

//export golua_idtointerface
func golua_idtointerface(id uint) interface{} {
	return id
}

//export golua_cfunctiontointerface
func golua_cfunctiontointerface(f *uintptr) interface{} {
	return f
}

//export golua_callallocf
func golua_callallocf(fp uintptr, ptr uintptr, osize uint, nsize uint) uintptr {
	return uintptr((*((*Alloc)(unsafe.Pointer(fp))))(unsafe.Pointer(ptr), osize, nsize))
}

//export go_panic_msghandler
func go_panic_msghandler(Li interface{}, z *C.char) {
	L := Li.(*State)
	s := C.GoString(z)

	panic(&LuaError{LUA_ERRERR, s, L.StackTrace()})
}
<|endoftext|>


<commit_before>package circularbuffer

import (
	"sync"
	"time"
)

// RateLimiter
type RateLimiter struct {
	sync.Mutex
	bag        map[string]*CircularBuffer
	maxHits    int
	timeWindow time.Duration
}

// NewRateLimiter returns a new initialized RateLimitter with maxHits is
// the maximal number of hits per time.Duration d.
func NewRateLimiter(maxHits int, d, cleanInterval time.Duration, quit chan struct{}) *RateLimiter {
	rl := &RateLimiter{
		bag:        make(map[string]*CircularBuffer),
		maxHits:    maxHits,
		timeWindow: d,
	}
	go rl.startCleanerDaemon(cleanInterval, quit)
	return rl
}

// Allow tries to add s to a circularbuffer and returns true if we have
// a free bucket, if not it will return false, which means ratelimit.
func (rl *RateLimiter) Allow(s string) bool {
	rl.Lock()
	if rl.bag[s] == nil {
		rl.bag[s] = NewCircularBuffer(rl.maxHits, rl.timeWindow)
	}
	res := rl.bag[s].Add(time.Now())
	rl.Unlock()
	return res
}

// DeleteOld removes old entries from state bag
func (rl *RateLimiter) DeleteOld() {
	rl.Lock()
	for k, cb := range rl.bag {
		if !cb.InUse() {
			delete(rl.bag, k)
		}
	}
	rl.Unlock()
}

func (rl *RateLimiter) startCleanerDaemon(d time.Duration, quit chan struct{}) {
	for {
		select {
		case <-quit:
			return
		case <-time.After(d):
			rl.DeleteOld()
		}
	}
}
<commit_msg>read and write locks in Allow() to select critical sections<commit_after>package circularbuffer

import (
	"sync"
	"time"
)

// RateLimiter
type RateLimiter struct {
	sync.RWMutex
	bag        map[string]*CircularBuffer
	maxHits    int
	timeWindow time.Duration
}

// NewRateLimiter returns a new initialized RateLimitter with maxHits is
// the maximal number of hits per time.Duration d.
func NewRateLimiter(maxHits int, d, cleanInterval time.Duration, quit chan struct{}) *RateLimiter {
	rl := &RateLimiter{
		bag:        make(map[string]*CircularBuffer),
		maxHits:    maxHits,
		timeWindow: d,
	}
	go rl.startCleanerDaemon(cleanInterval, quit)
	return rl
}

// Allow tries to add s to a circularbuffer and returns true if we have
// a free bucket, if not it will return false, which means ratelimit.
func (rl *RateLimiter) Allow(s string) bool {
	var source *CircularBuffer
	var present bool

	rl.RLock()
	if source, present = rl.bag[s]; !present {
		rl.RUnlock()
		rl.Lock()
		source = NewCircularBuffer(rl.maxHits, rl.timeWindow)
		rl.bag[s] = source
		rl.Unlock()
	} else {
		rl.RUnlock()
	}
	present = source.Add(time.Now())
	return present
}

// DeleteOld removes old entries from state bag
func (rl *RateLimiter) DeleteOld() {
	rl.Lock()
	for k, cb := range rl.bag {
		if !cb.InUse() {
			delete(rl.bag, k)
		}
	}
	rl.Unlock()
}

func (rl *RateLimiter) startCleanerDaemon(d time.Duration, quit chan struct{}) {
	for {
		select {
		case <-quit:
			return
		case <-time.After(d):
			rl.DeleteOld()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"github.com/cloudfoundry/cli/cf/api"
	"github.com/cloudfoundry/cli/cf/app"
	"github.com/cloudfoundry/cli/cf/command_factory"
	"github.com/cloudfoundry/cli/cf/command_runner"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/i18n"
	"github.com/cloudfoundry/cli/cf/manifest"
	"github.com/cloudfoundry/cli/cf/net"
	"github.com/cloudfoundry/cli/cf/panic_printer"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/cf/terminal"
	"github.com/cloudfoundry/cli/cf/trace"
	"github.com/codegangsta/cli"
)

var deps = setupDependencies()

type cliDependencies struct {
	termUI         terminal.UI
	configRepo     configuration.Repository
	manifestRepo   manifest.ManifestRepository
	apiRepoLocator api.RepositoryLocator
	gateways       map[string]net.Gateway
}

func setupDependencies() (deps *cliDependencies) {
	deps = new(cliDependencies)

	deps.termUI = terminal.NewUI(os.Stdin)

	deps.manifestRepo = manifest.NewManifestDiskRepository()

	deps.configRepo = configuration.NewRepositoryFromFilepath(configuration.DefaultFilePath(), func(err error) {
		if err != nil {
			deps.termUI.Failed(fmt.Sprintf("Config error: %s", err))
		}
	})

	i18n.T = i18n.Init(deps.configRepo)

	terminal.UserAskedForColors = deps.configRepo.ColorEnabled()
	terminal.InitColorSupport()

	if os.Getenv("CF_TRACE") != "" {
		trace.Logger = trace.NewLogger(os.Getenv("CF_TRACE"))
	} else {
		trace.Logger = trace.NewLogger(deps.configRepo.Trace())
	}

	deps.gateways = map[string]net.Gateway{
		"auth":             net.NewUAAGateway(deps.configRepo),
		"cloud-controller": net.NewCloudControllerGateway(deps.configRepo, time.Now),
		"uaa":              net.NewUAAGateway(deps.configRepo),
	}
	deps.apiRepoLocator = api.NewRepositoryLocator(deps.configRepo, deps.gateways)

	return
}

func main() {
	defer handlePanics()
	defer deps.configRepo.Close()

	cmdFactory := command_factory.NewFactory(deps.termUI, deps.configRepo, deps.manifestRepo, deps.apiRepoLocator)
	requirementsFactory := requirements.NewFactory(deps.termUI, deps.configRepo, deps.apiRepoLocator)
	cmdRunner := command_runner.NewRunner(cmdFactory, requirementsFactory)

	theApp := app.NewApp(cmdRunner, cmdFactory.CommandMetadatas()...)
	//command `cf` without argument
	if len(os.Args) == 1 {
		theApp.Run(os.Args)
	} else if cmdFactory.CheckIfCoreCmdExists(os.Args[1]) {
		callCoreCommand(os.Args[1:], theApp)
	} else {
		location := deps.configRepo.Plugins()[os.Args[1]]
		if location != "" {
			callPlugin(os.Args[1:], location)
		} else {
			theApp.Run(os.Args)
		}
	}
}

func gatewaySliceFromMap(gateway_map map[string]net.Gateway) []net.WarningProducer {
	gateways := []net.WarningProducer{}
	for _, gateway := range gateway_map {
		gateways = append(gateways, gateway)
	}
	return gateways
}

func init() {
	cli.CommandHelpTemplate = `NAME:
   {{.Name}} - {{.Description}}
{{with .ShortName}}
ALIAS:
   {{.}}
{{end}}
USAGE:
   {{.Usage}}{{with .Flags}}

OPTIONS:
{{range .}}   {{.}}
{{end}}{{else}}
{{end}}`

}

func handlePanics() {
	panic_printer.UI = terminal.NewUI(os.Stdin)

	commandArgs := strings.Join(os.Args, " ")
	stackTrace := generateBacktrace()

	err := recover()
	panic_printer.DisplayCrashDialog(err, commandArgs, stackTrace)

	if err != nil {
		os.Exit(1)
	}
}

func generateBacktrace() string {
	stackByteCount := 0
	STACK_SIZE_LIMIT := 1024 * 1024
	var bytes []byte
	for stackSize := 1024; (stackByteCount == 0 || stackByteCount == stackSize) && stackSize < STACK_SIZE_LIMIT; stackSize = 2 * stackSize {
		bytes = make([]byte, stackSize)
		stackByteCount = runtime.Stack(bytes, true)
	}
	stackTrace := "\t" + strings.Replace(string(bytes), "\n", "\n\t", -1)
	return stackTrace
}

func callCoreCommand(args []string, theApp *cli.App) {
	err := theApp.Run(args)
	if err != nil {
		os.Exit(1)
	}
	gateways := gatewaySliceFromMap(deps.gateways)

	warningsCollector := net.NewWarningsCollector(deps.termUI, gateways...)
	warningsCollector.PrintWarnings()
}

func callPlugin(args []string, location string) {
	cmd := exec.Command("go", "run", location)
	cmd.Stdout = os.Stdout
	err := cmd.Start()
	if err != nil {
		fmt.Println("Error running plugin command: ", err)
		os.Exit(1)
	}
	cmd.Wait()
}
<commit_msg>fix argument in callCoreCommand()<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"github.com/cloudfoundry/cli/cf/api"
	"github.com/cloudfoundry/cli/cf/app"
	"github.com/cloudfoundry/cli/cf/command_factory"
	"github.com/cloudfoundry/cli/cf/command_runner"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/i18n"
	"github.com/cloudfoundry/cli/cf/manifest"
	"github.com/cloudfoundry/cli/cf/net"
	"github.com/cloudfoundry/cli/cf/panic_printer"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/cf/terminal"
	"github.com/cloudfoundry/cli/cf/trace"
	"github.com/codegangsta/cli"
)

var deps = setupDependencies()

type cliDependencies struct {
	termUI         terminal.UI
	configRepo     configuration.Repository
	manifestRepo   manifest.ManifestRepository
	apiRepoLocator api.RepositoryLocator
	gateways       map[string]net.Gateway
}

func setupDependencies() (deps *cliDependencies) {
	deps = new(cliDependencies)

	deps.termUI = terminal.NewUI(os.Stdin)

	deps.manifestRepo = manifest.NewManifestDiskRepository()

	deps.configRepo = configuration.NewRepositoryFromFilepath(configuration.DefaultFilePath(), func(err error) {
		if err != nil {
			deps.termUI.Failed(fmt.Sprintf("Config error: %s", err))
		}
	})

	i18n.T = i18n.Init(deps.configRepo)

	terminal.UserAskedForColors = deps.configRepo.ColorEnabled()
	terminal.InitColorSupport()

	if os.Getenv("CF_TRACE") != "" {
		trace.Logger = trace.NewLogger(os.Getenv("CF_TRACE"))
	} else {
		trace.Logger = trace.NewLogger(deps.configRepo.Trace())
	}

	deps.gateways = map[string]net.Gateway{
		"auth":             net.NewUAAGateway(deps.configRepo),
		"cloud-controller": net.NewCloudControllerGateway(deps.configRepo, time.Now),
		"uaa":              net.NewUAAGateway(deps.configRepo),
	}
	deps.apiRepoLocator = api.NewRepositoryLocator(deps.configRepo, deps.gateways)

	return
}

func main() {
	defer handlePanics()
	defer deps.configRepo.Close()

	cmdFactory := command_factory.NewFactory(deps.termUI, deps.configRepo, deps.manifestRepo, deps.apiRepoLocator)
	requirementsFactory := requirements.NewFactory(deps.termUI, deps.configRepo, deps.apiRepoLocator)
	cmdRunner := command_runner.NewRunner(cmdFactory, requirementsFactory)

	theApp := app.NewApp(cmdRunner, cmdFactory.CommandMetadatas()...)
	//command `cf` without argument
	if len(os.Args) == 1 {
		theApp.Run(os.Args)
	} else if cmdFactory.CheckIfCoreCmdExists(os.Args[1]) {
		callCoreCommand(os.Args[0:], theApp)
	} else {
		location := deps.configRepo.Plugins()[os.Args[1]]
		if location != "" {
			callPlugin(os.Args[1:], location)
		} else {
			theApp.Run(os.Args)
		}
	}
}

func gatewaySliceFromMap(gateway_map map[string]net.Gateway) []net.WarningProducer {
	gateways := []net.WarningProducer{}
	for _, gateway := range gateway_map {
		gateways = append(gateways, gateway)
	}
	return gateways
}

func init() {
	cli.CommandHelpTemplate = `NAME:
   {{.Name}} - {{.Description}}
{{with .ShortName}}
ALIAS:
   {{.}}
{{end}}
USAGE:
   {{.Usage}}{{with .Flags}}

OPTIONS:
{{range .}}   {{.}}
{{end}}{{else}}
{{end}}`

}

func handlePanics() {
	panic_printer.UI = terminal.NewUI(os.Stdin)

	commandArgs := strings.Join(os.Args, " ")
	stackTrace := generateBacktrace()

	err := recover()
	panic_printer.DisplayCrashDialog(err, commandArgs, stackTrace)

	if err != nil {
		os.Exit(1)
	}
}

func generateBacktrace() string {
	stackByteCount := 0
	STACK_SIZE_LIMIT := 1024 * 1024
	var bytes []byte
	for stackSize := 1024; (stackByteCount == 0 || stackByteCount == stackSize) && stackSize < STACK_SIZE_LIMIT; stackSize = 2 * stackSize {
		bytes = make([]byte, stackSize)
		stackByteCount = runtime.Stack(bytes, true)
	}
	stackTrace := "\t" + strings.Replace(string(bytes), "\n", "\n\t", -1)
	return stackTrace
}

func callCoreCommand(args []string, theApp *cli.App) {
	err := theApp.Run(args)
	if err != nil {
		os.Exit(1)
	}
	gateways := gatewaySliceFromMap(deps.gateways)

	warningsCollector := net.NewWarningsCollector(deps.termUI, gateways...)
	warningsCollector.PrintWarnings()
}

func callPlugin(args []string, location string) {
	cmd := exec.Command("go", "run", location)
	cmd.Stdout = os.Stdout
	err := cmd.Start()
	if err != nil {
		fmt.Println("Error running plugin command: ", err)
		os.Exit(1)
	}
	cmd.Wait()
}
<|endoftext|>