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

package boosd

import (
	"bytes"
	"fmt"
	"github.com/bpowers/boosd/runtime"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"strconv"
	"text/template"
	"unicode"
)

const fileTmpl = `
{{define "modelTmpl"}}
var (
	mdl{{$.CamelName}}Name = "{{$.Name}}"
	mdl{{$.CamelName}}Vars = map[string]runtime.Var{
{{range $.Vars}}
		"{{.Name}}": runtime.Var{"{{.Name}}", runtime.{{.Type}}},{{end}}
	}
)

type sim{{$.CamelName}} struct {
	runtime.BaseSim
}

type mdl{{$.CamelName}} struct {
	runtime.BaseModel
}

func sim{{$.CamelName}}Step(s *runtime.BaseSim, dt float64) {
{{range $.Equations}}
	{{.}}
{{end}}
{{range $.Stocks}}
	{{.}}
{{end}}
}

func (m *mdl{{$.CamelName}}) NewSim() runtime.Sim {
	ts := runtime.Timespec{
		Start:    {{$.Time.Start}},
		End:      {{$.Time.End}},
		DT:       {{$.Time.DT}},
		SaveStep: {{$.Time.SaveStep}},
	}
	tables := map[string]runtime.Table{}
	consts := runtime.Data{}

	s := new(sim{{$.CamelName}})
	s.Init(m, ts, tables, consts)
	s.Step = sim{{$.CamelName}}Step

{{range $.Initials}}
	{{.}}
{{end}}
	s.Curr["time"] = ts.Start

	runtime.RegisterSim(mdl{{$.CamelName}}Name, s)

	return s
}
{{end}}

package main

import (
	"github.com/bpowers/boosd/runtime"
)

{{range $.Models}}{{template "modelTmpl" .}}{{end}}

func init() {
	m := &mdlMain{
		runtime.BaseModel{
			MName: mdlMainName,
			Vars:  mdlMainVars,
		},
	}

	runtime.RegisterModel(m)
}

func main() {
	runtime.Main()
}
`

type genModel struct {
	Name      string
	CamelName string // camelcased
	Vars      map[string]runtime.Var
	Time      runtime.Timespec
	Equations []string
	Stocks    []string
	Initials  []string
}

type generator struct {
	Models map[string]*genModel
	curr *genModel
}

func (g *generator) declList(list []Decl) {
}

// constEval returns the float64 value represented by Expr, or an
// error if it can't be evaluated at compile time.
func constEval(e Expr) (v float64, err error) {
	val, ok := e.(*UnitExpr)
	if !ok {
		err = fmt.Errorf("timespec val %T not UnitExpr", e)
		return
	}
	basic, ok := val.X.(*BasicLit)
	if !ok {
		err = fmt.Errorf("timespec val %T not BasicLit", val.X)
		return
	}
	return strconv.ParseFloat(basic.Value, 64)
}

func isConst(e Expr) bool {
	if _, err := constEval(e); err == nil {
		return true
	}
	return false
}

func kvConvert(e Expr) (k string, v Expr, err error) {
	kv, ok := e.(*KeyValueExpr)
	if !ok {
		err = fmt.Errorf("e %T not KVeyValueExpr: %v", e, e)
		return
	}
	ident, ok := kv.Key.(*Ident)
	if !ok {
		err = fmt.Errorf("e key %T not Ident: %v", kv.Key, kv.Key)
		return
	}
	return ident.Name, kv.Value, nil
}

func (g *generator) timespec(elts []Expr) {
	for _, e := range elts {
		k, val, err := kvConvert(e)
		if err != nil {
			panic(fmt.Sprintf("kvConvert(%v): %s", e, err))
		}
		v, err := constEval(val)
		if err != nil {
			panic(fmt.Sprintf("timespec constEval(%v): %s",
				val, err))
		}
		switch k {
		case "start":
			g.curr.Time.Start = v
		case "end":
			g.curr.Time.End = v
		case "dt":
			g.curr.Time.DT = v
		case "save_step":
			g.curr.Time.SaveStep = v
		default:
			panic(fmt.Sprintf("timespec unknown key %s", k))
		}
	}
}

var identAux = &Ident{Name: "aux"}

func varFromDecl(d *VarDecl) (v runtime.Var, err error) {
	if d.Type == nil {
		d.Type = identAux
	}
	return runtime.Var{d.Name.Name, runtime.TypeForName(d.Type.Name)}, nil
}

func (g *generator) initial(name string, expr Expr) {
	val, err := constEval(expr)
	if err != nil {
		panic(fmt.Sprintf("initial(%s): non-const %v", name, expr))
	}
	init := fmt.Sprintf(`s.Curr["%s"] = %f`, name, val)
	g.curr.Initials = append(g.curr.Initials, init)
}

func (g *generator) stock(name string, expr Expr) {
	cl, ok := expr.(*CompositeLit)
	if !ok {
		panic(fmt.Sprintf("stock is %T, not CompositeLit", expr))
	}
	var in, out string
	for _, e := range cl.Elts {
		k, val, err := kvConvert(e)
		if err != nil {
			panic(fmt.Sprintf("kvConvert(%s): %s", name, err))
		}
		switch k {
		case "initial":
			g.initial(name, val)
		case "inflow":
			in = fmt.Sprintf("%s", val)
		case "outflow":
			out = fmt.Sprintf("-(%s)", val)
		default:
			panic(fmt.Sprintf("stock(%s): unknown k %s",
				name, k))
		}
	}
	eqn := fmt.Sprintf(`s.Next["%s"] = s.Curr["%s"] + (%s %s)*dt`, name, name, in, out)
	g.curr.Stocks = append(g.curr.Stocks, eqn)
}

func (g *generator) expr(name string, expr Expr) {
	if g.curr.Vars[name].Type == runtime.TyAux {
		if isConst(expr) {
			g.initial(name, expr)
			eqn := fmt.Sprintf(`s.Next["%s"] = %s`, name, expr)
			g.curr.Stocks = append(g.curr.Stocks, eqn)
		} else {
			eqn := fmt.Sprintf(`s.Curr["%s"] = %s`, name, expr)
			g.curr.Equations = append(g.curr.Equations, eqn)
		}
	} else {
		eqn := fmt.Sprintf(`s.Curr["%s"] = %s`, name, expr)
		g.curr.Equations = append(g.curr.Equations, eqn)
	}
}

func (g *generator) assign(s *AssignStmt) {
	if s.Lhs.Name.Name == "timespec" {
		c, ok := s.Rhs.(*CompositeLit)
		if !ok {
			panic(fmt.Sprintf("timespec is %T, not CompositeLit",
				s.Rhs))
		}
		g.timespec(c.Elts)
		return
	}
	v, err := varFromDecl(s.Lhs)
	if err != nil {
		panic(fmt.Sprintf("varFromDecl(%v): %s", s.Lhs, err))
	}
	g.curr.Vars[v.Name] = v
	if v.Type == runtime.TyStock {
		g.stock(v.Name, s.Rhs)
	} else {
		g.expr(v.Name, s.Rhs)
	}
}

func (g *generator) stmt(s Stmt) {
	switch ss := s.(type) {
	case *AssignStmt:
		g.assign(ss)
	default:
		log.Printf("stmt(%T): unimplemented - %v", s, s)
	}
}

func (g *generator) model(m *ModelDecl) error {
	name := m.Name.Name
	camelName := fmt.Sprintf("%c%s", unicode.ToUpper(rune(name[0])), name[1:])
	g.curr = &genModel{
		Name: name,
		CamelName: camelName,
		Vars: map[string]runtime.Var{},
		Equations: []string{},
		Stocks: []string{},
		Initials: []string{},
	}
	for _, s := range m.Body.List {
		g.stmt(s)
	}
	g.Models[m.Name.Name] = g.curr
	g.curr = nil

	return nil
}

func (g *generator) file(f *File) ([]byte, error) {
	for _, d := range f.Decls {
		md, ok := d.(*ModelDecl)
		if !ok {
			log.Printf("top level decl that isn't a model: %v (%T)", d, d)
			continue
		}
		if err := g.model(md); err != nil {
			return nil, fmt.Errorf("g.model: %s", err)
		}
	}

	var buf bytes.Buffer
	tmpl := template.New("model.go")
	if _, err := tmpl.Parse(fileTmpl); err != nil {
		panic(fmt.Sprintf("Parse(modelTmpl): %s", err))
	}
	if err := tmpl.Execute(&buf, g); err != nil {
		panic(fmt.Sprintf("Execute(%v): %s", g, err))
	}

	return buf.Bytes(), nil
}

func GenGo(f *File) (*ast.File, error) {
	g := &generator{
		Models: map[string]*genModel{},
	}

	code, err := g.file(f)
	if err != nil {
		return nil, fmt.Errorf("g.file: %s", err)
	}
	//log.Printf("c: %s", code)

	fset := token.NewFileSet()
	goFile, err := parser.ParseFile(fset, "model.go", code, 0)
	if err != nil {
		return nil, err
	}

	return goFile, nil
}
<commit_msg>lib: generate: support biflows<commit_after>// Copyright 2013 Bobby Powers. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package boosd

import (
	"bytes"
	"fmt"
	"github.com/bpowers/boosd/runtime"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"strconv"
	"text/template"
	"unicode"
)

const fileTmpl = `
{{define "modelTmpl"}}
var (
	mdl{{$.CamelName}}Name = "{{$.Name}}"
	mdl{{$.CamelName}}Vars = map[string]runtime.Var{
{{range $.Vars}}
		"{{.Name}}": runtime.Var{"{{.Name}}", runtime.{{.Type}}},{{end}}
	}
)

type sim{{$.CamelName}} struct {
	runtime.BaseSim
}

type mdl{{$.CamelName}} struct {
	runtime.BaseModel
}

func sim{{$.CamelName}}Step(s *runtime.BaseSim, dt float64) {
{{range $.Equations}}
	{{.}}
{{end}}
{{range $.Stocks}}
	{{.}}
{{end}}
}

func (m *mdl{{$.CamelName}}) NewSim() runtime.Sim {
	ts := runtime.Timespec{
		Start:    {{$.Time.Start}},
		End:      {{$.Time.End}},
		DT:       {{$.Time.DT}},
		SaveStep: {{$.Time.SaveStep}},
	}
	tables := map[string]runtime.Table{}
	consts := runtime.Data{}

	s := new(sim{{$.CamelName}})
	s.Init(m, ts, tables, consts)
	s.Step = sim{{$.CamelName}}Step

{{range $.Initials}}
	{{.}}
{{end}}
	s.Curr["time"] = ts.Start

	runtime.RegisterSim(mdl{{$.CamelName}}Name, s)

	return s
}
{{end}}

package main

import (
	"github.com/bpowers/boosd/runtime"
)

{{range $.Models}}{{template "modelTmpl" .}}{{end}}

func init() {
	m := &mdlMain{
		runtime.BaseModel{
			MName: mdlMainName,
			Vars:  mdlMainVars,
		},
	}

	runtime.RegisterModel(m)
}

func main() {
	runtime.Main()
}
`

type genModel struct {
	Name      string
	CamelName string // camelcased
	Vars      map[string]runtime.Var
	Time      runtime.Timespec
	Equations []string
	Stocks    []string
	Initials  []string
}

type generator struct {
	Models map[string]*genModel
	curr *genModel
}

func (g *generator) declList(list []Decl) {
}

// constEval returns the float64 value represented by Expr, or an
// error if it can't be evaluated at compile time.
func constEval(e Expr) (v float64, err error) {
	val, ok := e.(*UnitExpr)
	if !ok {
		err = fmt.Errorf("timespec val %T not UnitExpr", e)
		return
	}
	basic, ok := val.X.(*BasicLit)
	if !ok {
		err = fmt.Errorf("timespec val %T not BasicLit", val.X)
		return
	}
	return strconv.ParseFloat(basic.Value, 64)
}

func isConst(e Expr) bool {
	if _, err := constEval(e); err == nil {
		return true
	}
	return false
}

func kvConvert(e Expr) (k string, v Expr, err error) {
	kv, ok := e.(*KeyValueExpr)
	if !ok {
		err = fmt.Errorf("e %T not KVeyValueExpr: %v", e, e)
		return
	}
	ident, ok := kv.Key.(*Ident)
	if !ok {
		err = fmt.Errorf("e key %T not Ident: %v", kv.Key, kv.Key)
		return
	}
	return ident.Name, kv.Value, nil
}

func (g *generator) timespec(elts []Expr) {
	for _, e := range elts {
		k, val, err := kvConvert(e)
		if err != nil {
			panic(fmt.Sprintf("kvConvert(%v): %s", e, err))
		}
		v, err := constEval(val)
		if err != nil {
			panic(fmt.Sprintf("timespec constEval(%v): %s",
				val, err))
		}
		switch k {
		case "start":
			g.curr.Time.Start = v
		case "end":
			g.curr.Time.End = v
		case "dt":
			g.curr.Time.DT = v
		case "save_step":
			g.curr.Time.SaveStep = v
		default:
			panic(fmt.Sprintf("timespec unknown key %s", k))
		}
	}
}

var identAux = &Ident{Name: "aux"}

func varFromDecl(d *VarDecl) (v runtime.Var, err error) {
	if d.Type == nil {
		d.Type = identAux
	}
	return runtime.Var{d.Name.Name, runtime.TypeForName(d.Type.Name)}, nil
}

func (g *generator) initial(name string, expr Expr) {
	val, err := constEval(expr)
	if err != nil {
		panic(fmt.Sprintf("initial(%s): non-const %v", name, expr))
	}
	init := fmt.Sprintf(`s.Curr["%s"] = %f`, name, val)
	g.curr.Initials = append(g.curr.Initials, init)
}

func (g *generator) stock(name string, expr Expr) {
	cl, ok := expr.(*CompositeLit)
	if !ok {
		panic(fmt.Sprintf("stock is %T, not CompositeLit", expr))
	}
	var bi, in, out string
	for _, e := range cl.Elts {
		k, val, err := kvConvert(e)
		if err != nil {
			panic(fmt.Sprintf("kvConvert(%s): %s", name, err))
		}
		switch k {
		case "initial":
			g.initial(name, val)
		case "biflow":
			bi = fmt.Sprintf("%s", val)
		case "inflow":
			in = fmt.Sprintf("%s", val)
		case "outflow":
			out = fmt.Sprintf("-(%s)", val)
		default:
			panic(fmt.Sprintf("stock(%s): unknown k %s",
				name, k))
		}
	}
	eqn := fmt.Sprintf(`s.Next["%s"] = s.Curr["%s"] + (%s %s %s)*dt`, name, name, bi, in, out)
	g.curr.Stocks = append(g.curr.Stocks, eqn)
}

func (g *generator) expr(name string, expr Expr) {
	if g.curr.Vars[name].Type == runtime.TyAux {
		if isConst(expr) {
			g.initial(name, expr)
			eqn := fmt.Sprintf(`s.Next["%s"] = %s`, name, expr)
			g.curr.Stocks = append(g.curr.Stocks, eqn)
		} else {
			eqn := fmt.Sprintf(`s.Curr["%s"] = %s`, name, expr)
			g.curr.Equations = append(g.curr.Equations, eqn)
		}
	} else {
		eqn := fmt.Sprintf(`s.Curr["%s"] = %s`, name, expr)
		g.curr.Equations = append(g.curr.Equations, eqn)
	}
}

func (g *generator) assign(s *AssignStmt) {
	if s.Lhs.Name.Name == "timespec" {
		c, ok := s.Rhs.(*CompositeLit)
		if !ok {
			panic(fmt.Sprintf("timespec is %T, not CompositeLit",
				s.Rhs))
		}
		g.timespec(c.Elts)
		return
	}
	v, err := varFromDecl(s.Lhs)
	if err != nil {
		panic(fmt.Sprintf("varFromDecl(%v): %s", s.Lhs, err))
	}
	g.curr.Vars[v.Name] = v
	if v.Type == runtime.TyStock {
		g.stock(v.Name, s.Rhs)
	} else {
		g.expr(v.Name, s.Rhs)
	}
}

func (g *generator) stmt(s Stmt) {
	switch ss := s.(type) {
	case *AssignStmt:
		g.assign(ss)
	default:
		log.Printf("stmt(%T): unimplemented - %v", s, s)
	}
}

func (g *generator) model(m *ModelDecl) error {
	name := m.Name.Name
	camelName := fmt.Sprintf("%c%s", unicode.ToUpper(rune(name[0])), name[1:])
	g.curr = &genModel{
		Name: name,
		CamelName: camelName,
		Vars: map[string]runtime.Var{},
		Equations: []string{},
		Stocks: []string{},
		Initials: []string{},
	}
	for _, s := range m.Body.List {
		g.stmt(s)
	}
	g.Models[m.Name.Name] = g.curr
	g.curr = nil

	return nil
}

func (g *generator) file(f *File) ([]byte, error) {
	for _, d := range f.Decls {
		md, ok := d.(*ModelDecl)
		if !ok {
			log.Printf("top level decl that isn't a model: %v (%T)", d, d)
			continue
		}
		if err := g.model(md); err != nil {
			return nil, fmt.Errorf("g.model: %s", err)
		}
	}

	var buf bytes.Buffer
	tmpl := template.New("model.go")
	if _, err := tmpl.Parse(fileTmpl); err != nil {
		panic(fmt.Sprintf("Parse(modelTmpl): %s", err))
	}
	if err := tmpl.Execute(&buf, g); err != nil {
		panic(fmt.Sprintf("Execute(%v): %s", g, err))
	}

	return buf.Bytes(), nil
}

func GenGo(f *File) (*ast.File, error) {
	g := &generator{
		Models: map[string]*genModel{},
	}

	code, err := g.file(f)
	if err != nil {
		return nil, fmt.Errorf("g.file: %s", err)
	}
	//log.Printf("c: %s", code)

	fset := token.NewFileSet()
	goFile, err := parser.ParseFile(fset, "model.go", code, 0)
	if err != nil {
		return nil, err
	}

	return goFile, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/RackSec/srslog"
	"github.com/buger/jsonparser"
	"github.com/urfave/cli"
)

var papertrailprefix string

func init() {
	app.Commands = append(app.Commands,
		cli.Command{
			Name:  "papertrail",
			Usage: "forward logs to papertrail",
			Action: func(c *cli.Context) {

				papertrailprefix = c.String("prefix")
				consumer := NewConsumer(globalFlags.Brokers, globalFlags.Groupid)
				consumer.Init(globalFlags.Topic)

				log.Println(c.String("papertrail"))

				go listenforchecks()
				papertrail := make(chan *data, 20)
				defer close(papertrail)
				go func(messages <-chan *Message, papertrail chan<- *data) {
					for msg := range messages {
						jsondata := &data{start: time.Now(), offset: msg.Offset}
						jsondata.Ts, _ = jsonparser.GetUnsafeString(msg.Data, "time")
						jsondata.Msg, _ = jsonparser.GetString(msg.Data, "log")
						jsondata.Environment, _ = jsonparser.GetUnsafeString(msg.Data, "environment")
						jsondata.Level, _ = jsonparser.GetUnsafeString(msg.Data, "level")
						jsondata.Stream, _ = jsonparser.GetUnsafeString(msg.Data, "stream")
						jsondata.Service, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "labels", "app")
						jsondata.Host, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "host")
						jsondata.IPAddress, _ = jsonparser.GetUnsafeString(msg.Data, "ip_address")
						jsondata.ServerIP, _ = jsonparser.GetUnsafeString(msg.Data, "server_ip")
						jsondata.DockerImage, _ = jsonparser.GetUnsafeString(msg.Data, "docker_image")
						jsondata.ContainerName, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "pod_name")
						jsondata.ContainerID, _ = jsonparser.GetUnsafeString(msg.Data, "docker", "container_id")

						// NOTE: GetString does some allocations, which might cause some overhead.
						jsondata.Msg = strings.TrimSpace(jsondata.Msg)
						papertrail <- jsondata
					}
				}(consumer.Chan, papertrail)
				go Sender(papertrail, c.String("papertrail"), c.String("cert"))

				consumer.StartConsumingTopic()

				consumer.Wait()
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "papertrail",
					Usage:  "The papertrail address where the logs should be forwarded",
					EnvVar: "PAPERTRAIL_ADDRESS",
				},
				cli.StringFlag{
					Name:   "cert",
					Usage:  "Papertrail root certificate path",
					Value:  "./papertrail-bundle.pem",
					EnvVar: "PAPERTRAIL_CERT",
				},
				cli.StringFlag{
					Name:   "prefix",
					Usage:  "String that is prepended to the system name",
					EnvVar: "PAPERTRAIL_PREFIX",
				},
			},
		})
}

type data struct {
	Ts            string `json:"ts"`
	Msg           string `json:"msg"`
	Level         string `json:"level"`
	Stream        string `json:"stream"`
	Service       string `json:"service"`
	Environment   string `json:"environment"`
	Host          string `json:"host"`
	IPAddress     string `json:"ip_address"` // legacy
	ServerIP      string `json:"server_ip"`
	DockerImage   string `json:"docker_image"`
	ContainerName string `json:"container_name"`
	ContainerID   string `json:"container_id"`

	start  time.Time
	offset int64
}

// Sender receives valid entries from chan 'c' and uploads them into papertrail over an encrypted TCP connection.
func Sender(c <-chan *data, address, cert string) {
	w, err := srslog.DialWithTLSCertPath("tcp+tls", address, srslog.LOG_INFO, "kafkatopapertrail", cert)
	if err != nil {
		log.Fatal(err)
	}

	var timechan chan time.Duration
	if globalFlags.Verbose {
		timechan = make(chan time.Duration, 10)
		defer close(timechan)
		go func(timechan chan time.Duration) {
			var (
				eventcount int
				counter    time.Duration
				ticker     = time.NewTicker(time.Second * 30)
			)

		loop:
			for {
				select {
				case duration, ok := <-timechan:
					if ok {
						counter += duration
						eventcount++
					} else {
						break loop
					}
				case <-ticker.C:
					if eventcount != 0 {
						counter = counter / time.Duration(eventcount)
					}
					log.Printf("Sent %d messages during the last 30 seconds, averaging %s per message\n", eventcount, counter)
					eventcount = 0
					counter = 0
				}
			}
		}(timechan)
	}

	w.SetFormatter(LogFormatter)
	for msg := range c {
		switch msg.Level {
		case "INFO":
			err = w.Info(msg.String())
		case "ALERT":
			err = w.Alert(msg.String())
		case "CRIT", "CRITICAL":
			err = w.Crit(msg.String())
		case "DEBUG":
			err = w.Debug(msg.String())
		case "EMERG":
			err = w.Emerg(msg.String())
		case "ERR", "ERROR":
			err = w.Err(msg.String())
		case "NOTICE":
			err = w.Notice(msg.String())
		case "WARNING", "WARN":
			err = w.Warning(msg.String())
		default:
			err = w.Info(msg.String())
			//      log.Printf("Unknown log level: %s\n", msg.Level)
		}
		if err != nil {
			log.Println(err)
		}
		if globalFlags.Verbose {
			timechan <- time.Since(msg.start)
		}
	}
}

func (d *data) String() string {
	prefix := papertrailprefix
	ts, _ := strconv.ParseInt(d.Ts, 10, 64)
	timestamp := time.Unix(ts, 0).Format(time.RFC3339)
	if d.Environment == "stg" {
		// A bit conflicting with the whole point of shared prefix but mostly a easy solution until some things are sorted.
		prefix = "staging-"
	}
	return fmt.Sprintf("%s|%s%s|%s|%s %s", timestamp, prefix, d.Service, d.ContainerName, timestamp, d.Msg)
}

func parseImage(image string) (repository, tag string) {
	// Docker image format is:
	// {registry}/(_|/r/{user_or_org})/{repository}:{tag}
	// Here we get the repository and tag from that format
	var (
		parts []string
	)
	parts = strings.Split(image, "/")
	image = parts[len(parts)-1]
	parts = strings.Split(image, ":")
	repository = parts[0]
	if len(parts) > 1 {
		tag = parts[1]
		if len(tag) > 8 {
			tag = tag[:8]
		}
	}
	return
}

// LogFormatter is a custom syslog formatter that uses log data to fill timestamp, hostname and tag instead of using local system information.
func LogFormatter(p srslog.Priority, hostname, tag, content string) (msg string) {
	parts := strings.SplitN(content, "|", 4)
	if len(parts) == 4 {
		msg = fmt.Sprintf("<%d>1 %s %s %s - - - %s", p, parts[0], parts[1], parts[2], parts[3])
	} else {
		msg = srslog.DefaultFormatter(p, hostname, tag, content)
	}
	return
}
<commit_msg>Proper environment label location<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/RackSec/srslog"
	"github.com/buger/jsonparser"
	"github.com/urfave/cli"
)

var papertrailprefix string

func init() {
	app.Commands = append(app.Commands,
		cli.Command{
			Name:  "papertrail",
			Usage: "forward logs to papertrail",
			Action: func(c *cli.Context) {

				papertrailprefix = c.String("prefix")
				consumer := NewConsumer(globalFlags.Brokers, globalFlags.Groupid)
				consumer.Init(globalFlags.Topic)

				log.Println(c.String("papertrail"))

				go listenforchecks()
				papertrail := make(chan *data, 20)
				defer close(papertrail)
				go func(messages <-chan *Message, papertrail chan<- *data) {
					for msg := range messages {
						jsondata := &data{start: time.Now(), offset: msg.Offset}
						jsondata.Ts, _ = jsonparser.GetUnsafeString(msg.Data, "time")
						jsondata.Msg, _ = jsonparser.GetString(msg.Data, "log")
						jsondata.Level, _ = jsonparser.GetUnsafeString(msg.Data, "level")
						jsondata.Stream, _ = jsonparser.GetUnsafeString(msg.Data, "stream")
						jsondata.Service, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "labels", "app")
						jsondata.Environment, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "labels", "environment")
						jsondata.Host, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "host")
						jsondata.IPAddress, _ = jsonparser.GetUnsafeString(msg.Data, "ip_address")
						jsondata.ServerIP, _ = jsonparser.GetUnsafeString(msg.Data, "server_ip")
						jsondata.DockerImage, _ = jsonparser.GetUnsafeString(msg.Data, "docker_image")
						jsondata.ContainerName, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "pod_name")
						jsondata.ContainerID, _ = jsonparser.GetUnsafeString(msg.Data, "docker", "container_id")

						// NOTE: GetString does some allocations, which might cause some overhead.
						jsondata.Msg = strings.TrimSpace(jsondata.Msg)
						papertrail <- jsondata
					}
				}(consumer.Chan, papertrail)
				go Sender(papertrail, c.String("papertrail"), c.String("cert"))

				consumer.StartConsumingTopic()

				consumer.Wait()
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "papertrail",
					Usage:  "The papertrail address where the logs should be forwarded",
					EnvVar: "PAPERTRAIL_ADDRESS",
				},
				cli.StringFlag{
					Name:   "cert",
					Usage:  "Papertrail root certificate path",
					Value:  "./papertrail-bundle.pem",
					EnvVar: "PAPERTRAIL_CERT",
				},
				cli.StringFlag{
					Name:   "prefix",
					Usage:  "String that is prepended to the system name",
					EnvVar: "PAPERTRAIL_PREFIX",
				},
			},
		})
}

type data struct {
	Ts            string `json:"ts"`
	Msg           string `json:"msg"`
	Level         string `json:"level"`
	Stream        string `json:"stream"`
	Service       string `json:"service"`
	Environment   string `json:"environment"`
	Host          string `json:"host"`
	IPAddress     string `json:"ip_address"` // legacy
	ServerIP      string `json:"server_ip"`
	DockerImage   string `json:"docker_image"`
	ContainerName string `json:"container_name"`
	ContainerID   string `json:"container_id"`

	start  time.Time
	offset int64
}

// Sender receives valid entries from chan 'c' and uploads them into papertrail over an encrypted TCP connection.
func Sender(c <-chan *data, address, cert string) {
	w, err := srslog.DialWithTLSCertPath("tcp+tls", address, srslog.LOG_INFO, "kafkatopapertrail", cert)
	if err != nil {
		log.Fatal(err)
	}

	var timechan chan time.Duration
	if globalFlags.Verbose {
		timechan = make(chan time.Duration, 10)
		defer close(timechan)
		go func(timechan chan time.Duration) {
			var (
				eventcount int
				counter    time.Duration
				ticker     = time.NewTicker(time.Second * 30)
			)

		loop:
			for {
				select {
				case duration, ok := <-timechan:
					if ok {
						counter += duration
						eventcount++
					} else {
						break loop
					}
				case <-ticker.C:
					if eventcount != 0 {
						counter = counter / time.Duration(eventcount)
					}
					log.Printf("Sent %d messages during the last 30 seconds, averaging %s per message\n", eventcount, counter)
					eventcount = 0
					counter = 0
				}
			}
		}(timechan)
	}

	w.SetFormatter(LogFormatter)
	for msg := range c {
		switch msg.Level {
		case "INFO":
			err = w.Info(msg.String())
		case "ALERT":
			err = w.Alert(msg.String())
		case "CRIT", "CRITICAL":
			err = w.Crit(msg.String())
		case "DEBUG":
			err = w.Debug(msg.String())
		case "EMERG":
			err = w.Emerg(msg.String())
		case "ERR", "ERROR":
			err = w.Err(msg.String())
		case "NOTICE":
			err = w.Notice(msg.String())
		case "WARNING", "WARN":
			err = w.Warning(msg.String())
		default:
			err = w.Info(msg.String())
			//      log.Printf("Unknown log level: %s\n", msg.Level)
		}
		if err != nil {
			log.Println(err)
		}
		if globalFlags.Verbose {
			timechan <- time.Since(msg.start)
		}
	}
}

func (d *data) String() string {
	prefix := papertrailprefix
	ts, _ := strconv.ParseInt(d.Ts, 10, 64)
	timestamp := time.Unix(ts, 0).Format(time.RFC3339)
	if d.Environment == "stg" {
		// A bit conflicting with the whole point of shared prefix but mostly a easy solution until some things are sorted.
		prefix = "staging-"
	}
	return fmt.Sprintf("%s|%s%s|%s|%s %s", timestamp, prefix, d.Service, d.ContainerName, timestamp, d.Msg)
}

func parseImage(image string) (repository, tag string) {
	// Docker image format is:
	// {registry}/(_|/r/{user_or_org})/{repository}:{tag}
	// Here we get the repository and tag from that format
	var (
		parts []string
	)
	parts = strings.Split(image, "/")
	image = parts[len(parts)-1]
	parts = strings.Split(image, ":")
	repository = parts[0]
	if len(parts) > 1 {
		tag = parts[1]
		if len(tag) > 8 {
			tag = tag[:8]
		}
	}
	return
}

// LogFormatter is a custom syslog formatter that uses log data to fill timestamp, hostname and tag instead of using local system information.
func LogFormatter(p srslog.Priority, hostname, tag, content string) (msg string) {
	parts := strings.SplitN(content, "|", 4)
	if len(parts) == 4 {
		msg = fmt.Sprintf("<%d>1 %s %s %s - - - %s", p, parts[0], parts[1], parts[2], parts[3])
	} else {
		msg = srslog.DefaultFormatter(p, hostname, tag, content)
	}
	return
}
<|endoftext|>


<commit_before>package bsw

import (
	"errors"
	"github.com/miekg/dns"
)

// Search for a possible wild card host by attempting to
// get an A record youmustconstructmoreplylons.[domain].
func GetWildCard(domain, serverAddr string) string {
	var fqdn = "youmustconstructmorepylons." + domain
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeA)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return ""
	}
	if len(in.Answer) < 1 {
		return ""
	}
	if a, ok := in.Answer[0].(*dns.A); ok {
		return a.A.String()
	} else {
		return ""
	}
}

// Search for a possible wild card host by attempting to
// get an AAAA record youmustconstructmoreplylons.[domain].
func GetWildCard6(domain, serverAddr string) string {
	var fqdn = "youmustconstructmorepylons." + domain
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeAAAA)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return ""
	}
	if len(in.Answer) < 1 {
		return ""
	}
	if a, ok := in.Answer[0].(*dns.AAAA); ok {
		return a.AAAA.String()
	} else {
		return ""
	}
}

// Attempt to get an A and CNAME record for a sub domain of domain.
func Dictionary(domain, subname, blacklist, serverAddr string) (Results, error) {
	results := Results{}
	var fqdn = subname + "." + domain
	ip, err := LookupName(fqdn, serverAddr)
	if err != nil {
		cfqdn, err := LookupCname(fqdn, serverAddr)
		if err != nil {
			return results, err
		}
		ip, err = LookupName(cfqdn, serverAddr)
		if err != nil {
			return results, err
		}
		if ip == blacklist {
			return results, errors.New("Returned IP in blackslist")
		}
		results = append(results, Result{Source: "Dictionary-CNAME", IP: ip, Hostname: fqdn}, Result{Source: "Dictionary-CNAME", IP: ip, Hostname: cfqdn})
		return results, nil
	}
	if ip == blacklist {
		return results, errors.New("Returned IP in blacklist")
	}
	results = append(results, Result{Source: "Dictionary", IP: ip, Hostname: fqdn})
	return results, nil
}

// Attempt to get an AAAA record for a sub domain of a domain.
func Dictionary6(domain, subname, blacklist, serverAddr string) (Results, error) {
	results := Results{}
	var fqdn = subname + "." + domain
	ip, err := LookupName6(fqdn, serverAddr)
	if err != nil {
		return results, err
	}
	if ip == blacklist {
		return results, errors.New("Returned IP in blacklist")
	}
	results = append(results, Result{Source: "Dictionary IPv6", IP: ip, Hostname: fqdn})
	return results, nil
}
<commit_msg>Refactor wild card functions to use generic.go<commit_after>package bsw

import (
	"errors"
)

const wildcardsub = "youmustcontstuctmoreplyons."

// Search for a possible wild card host by attempting to
// get an A record for wildcardsub + domain.
func GetWildCard(domain, serverAddr string) string {
	var fqdn = wildcardsub + domain
	ip, _ := LookupName(fqdn, serverAddr)
	return ip
}

// Search for a possible wild card host by attempting to
// get an AAAA record wildcardsub + domain.
func GetWildCard6(domain, serverAddr string) string {
	var fqdn = wildcardsub + domain
        ip, _ := LookupName6(fqdn, serverAddr)
        return ip
}

// Attempt to get an A and CNAME record for a sub domain of domain.
func Dictionary(domain, subname, blacklist, serverAddr string) (Results, error) {
	results := Results{}
	var fqdn = subname + "." + domain
	ip, err := LookupName(fqdn, serverAddr)
	if err != nil {
		cfqdn, err := LookupCname(fqdn, serverAddr)
		if err != nil {
			return results, err
		}
		ip, err = LookupName(cfqdn, serverAddr)
		if err != nil {
			return results, err
		}
		if ip == blacklist {
			return results, errors.New("Returned IP in blackslist")
		}
		results = append(results, Result{Source: "Dictionary-CNAME", IP: ip, Hostname: fqdn}, Result{Source: "Dictionary-CNAME", IP: ip, Hostname: cfqdn})
		return results, nil
	}
	if ip == blacklist {
		return results, errors.New("Returned IP in blacklist")
	}
	results = append(results, Result{Source: "Dictionary", IP: ip, Hostname: fqdn})
	return results, nil
}

// Attempt to get an AAAA record for a sub domain of a domain.
func Dictionary6(domain, subname, blacklist, serverAddr string) (Results, error) {
	results := Results{}
	var fqdn = subname + "." + domain
	ip, err := LookupName6(fqdn, serverAddr)
	if err != nil {
		return results, err
	}
	if ip == blacklist {
		return results, errors.New("Returned IP in blacklist")
	}
	results = append(results, Result{Source: "Dictionary IPv6", IP: ip, Hostname: fqdn})
	return results, nil
}
<|endoftext|>


<commit_before>package phraseapp

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"

	"github.com/bgentry/speakeasy"
)

type AuthHandler struct {
	Username string `cli:"opt --username desc='username used for authentication'"`
	Token    string `cli:"opt --token desc='token used for authentication'"`
	TFA      bool   `cli:"opt --tfa desc='use Two-Factor Authentication'"`
	Host     string "https://api.phraseapp.com/"
	Config   string `cli:"opt --path default=$HOME/.config/phraseapp/config.json desc='path to the config file'"`
}

var authH *AuthHandler

func RegisterAuthHandler(a *AuthHandler) {
	authH = a
}

func (a *AuthHandler) readConfig() error {
	tmpA := new(AuthHandler)

	path := os.ExpandEnv(a.Config)
	_, err := os.Stat(path)
	switch {
	case err == nil: // ignore
		fh, err := os.Open(path)
		if err != nil {
			return err
		}
		defer fh.Close()
		err = json.NewDecoder(fh).Decode(&tmpA)
		if err != nil {
			return err
		}
	case os.IsNotExist(err):
		// ignore
	default:
		return err
	}

	// Only set token if username not specified on commandline.
	if tmpA.Token != "" && a.Token == "" && a.Username == "" {
		a.Token = tmpA.Token
	} else if tmpA.Username != "" && a.Username == "" {
		a.Username = tmpA.Username
	}

	if tmpA.TFA && a.Username == "" {
		a.TFA = tmpA.TFA
	}

	// Set custom host if specified
	if tmpA.Host != "" {
		a.Host = tmpA.Host
	}else{
		a.Host = "https://api.phraseapp.com"
	}

	return nil
}

func authenticate(req *http.Request) error {
	if authH == nil {
		return fmt.Errorf("no auth handler registered")
	}

	if err := authH.readConfig(); err != nil {
		return err
	}

	if err := authH.validate(); err != nil {
		return err
	}

	switch {
	case authH.Token != "":
		req.Header.Set("Authorization", "token "+authH.Token)
	case authH.Username != "":
		pwd, err := speakeasy.Ask("Password: ")
		if err != nil {
			return err
		}
		req.SetBasicAuth(authH.Username, pwd)

		if authH.TFA { // TFA only required for username+password based login.
			token, err := speakeasy.Ask("TFA-Token: ")
			if err != nil {
				return err
			}
			req.Header.Set("X-PhraseApp-OTP", token)
		}
	}

	return nil
}

func (ah *AuthHandler) validate() error {
	switch {
	case ah.Username == "" && ah.Token == "":
		return fmt.Errorf("either username or token must be given")
	default:
		return nil
	}
}
<commit_msg>initial, factored out config logic<commit_after>package phraseapp

import (
	"fmt"
	"net/http"

	"github.com/bgentry/speakeasy"
)

type DefaultParams Action
type Action map[string]map[string]interface{}

type AuthCredentials struct {
	Username string `cli:"opt --username desc='username used for authentication'"`
	Token    string `cli:"opt --token desc='token used for authentication'"`
	TFA      bool   `cli:"opt --tfa desc='use Two-Factor Authentication'"`
	Host     string `cli:"opt --api-host desc='api-host to use'`
	Config   string `cli:"opt --path default=$HOME/.config/phraseapp/config.json desc='path to the config file'"`
}

var authC *AuthCredentials

func RegisterAuthCredentials(cmdAuth *AuthCredentials, defaultCredentials *AuthCredentials) {
	if authC == nil {
		authC = new(AuthCredentials)
	}

	if cmdAuth.Token != "" && authC.Token == "" && authC.Username == "" {
		authC.Token = cmdAuth.Token
	} else if cmdAuth.Username != "" && authC.Username == "" {
		authC.Username = cmdAuth.Username
	}

	if cmdAuth.TFA && authC.Username == "" {
		authC.TFA = cmdAuth.TFA || defaultCredentials.TFA
	}

	notSet := authC.Token == "" && authC.Username == ""
	if notSet && defaultCredentials.Token != "" {
		authC.Token = defaultCredentials.Token
	}
	if notSet && defaultCredentials.Username != "" {
		authC.Username = defaultCredentials.Username
	}
}

func authenticate(req *http.Request) error {
	if authC == nil {
		return fmt.Errorf("no auth handler registered")
	}

	if err := authC.validate(); err != nil {
		return err
	}

	switch {
	case authC.Token != "":
		req.Header.Set("Authorization", "token "+authC.Token)
	case authC.Username != "":
		pwd, err := speakeasy.Ask("Password: ")
		if err != nil {
			return err
		}
		req.SetBasicAuth(authC.Username, pwd)

		if authC.TFA { // TFA only required for username+password based login.
			token, err := speakeasy.Ask("TFA-Token: ")
			if err != nil {
				return err
			}
			req.Header.Set("X-PhraseApp-OTP", token)
		}
	}

	return nil
}

func (ah *AuthCredentials) validate() error {
	switch {
	case ah.Username == "" && ah.Token == "":
		return fmt.Errorf("either username or token must be given")
	default:
		return nil
	}
}
<|endoftext|>


<commit_before>package entity

import (
	"fmt"

	"rsprd.com/spread/pkg/deploy"
	"rsprd.com/spread/pkg/image"

	kube "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/validation"
)

var DefaultPodSpec = kube.PodSpec{
	RestartPolicy: kube.RestartPolicyAlways,
	DNSPolicy:     kube.DNSDefault,
}

// Pod represents kube.Pod in the Redspread hierarchy.
type Pod struct {
	base
	pod        *kube.Pod
	containers []*Container
}

func NewPod(kubePod *kube.Pod, defaults kube.ObjectMeta, source string, objects ...deploy.KubeObject) (*Pod, error) {
	if kubePod == nil {
		return nil, fmt.Errorf("cannot create Pod from nil `%s`", source)
	}

	// deep copy
	kubePod, err := copyPod(kubePod)
	if err != nil {
		return nil, err
	}

	base, err := newBase(EntityPod, defaults, source, objects)
	if err != nil {
		return nil, fmt.Errorf("could not create Pod from `%s`: %v", source, err)
	}

	pod := Pod{base: base}
	for _, v := range kubePod.Spec.Containers {
		container, err := NewContainer(v, defaults, source)
		if err != nil {
			return nil, err
		} else {
			pod.containers = append(pod.containers, container)
		}
	}
	kubePod.Spec.Containers = []kube.Container{}

	base.setDefaults(kubePod)
	if err := validatePod(kubePod, true); err != nil {
		return nil, fmt.Errorf("could not create Pod from `%s`: %v", source, err)
	}

	pod.pod = kubePod
	return &pod, nil
}

func NewPodFromPodSpec(meta kube.ObjectMeta, podSpec kube.PodSpec, defaults kube.ObjectMeta, source string, objects ...deploy.KubeObject) (*Pod, error) {
	pod := kube.Pod{
		ObjectMeta: meta,
		Spec:       podSpec,
	}
	return NewPod(&pod, defaults, source, objects...)
}

func newDefaultPod(meta kube.ObjectMeta, source string) (*Pod, error) {
	return NewPodFromPodSpec(meta, DefaultPodSpec, kube.ObjectMeta{}, source)
}

func (c *Pod) Deployment() (*deploy.Deployment, error) {
	return nil, nil
}

func (c *Pod) Images() (images []*image.Image) {
	for _, v := range c.containers {
		images = append(images, v.Images()...)
	}
	return
}

func (c *Pod) Attach(curEntity Entity) error {
	if err := c.validAttach(curEntity); err != nil {
		return err
	}
	for {
		switch e := curEntity.(type) {
		case *Image:
			container, err := newDefaultContainer(e.name(), e.Source())
			if err != nil {
				return err
			}

			err = container.Attach(e)
			curEntity = container
			break
		case *Container:
			c.containers = append(c.containers, e)
			return nil
		}
	}
}

func (c *Pod) name() string {
	return c.pod.ObjectMeta.Name
}

func (c *Pod) children() (children []Entity) {
	for _, v := range c.containers {
		children = append(children, v)
	}
	return
}

func (c *Pod) data() (pod *kube.Pod, objects deploy.Deployment, err error) {
	containers := []kube.Container{}
	for _, container := range c.containers {
		kubeContainer, cObj, err := container.data()
		if err != nil {
			return nil, objects, err
		}
		containers = append(containers, kubeContainer)
		// add containers objects
		objects.AddDeployment(cObj)
	}

	if len(containers) == 0 {
		return nil, objects, ErrorEntityNotReady
	}

	// add own objects
	objects.AddDeployment(c.objects)

	pod, err = copyPod(c.pod)
	if err != nil {
		return nil, objects, err
	}

	pod.Spec.Containers = containers
	return pod, objects, nil
}

func copyPod(pod *kube.Pod) (*kube.Pod, error) {
	copy, err := kube.Scheme.DeepCopy(pod)
	if err != nil {
		return nil, err
	}

	return copy.(*kube.Pod), nil
}

func validatePod(pod *kube.Pod, ignoreContainers bool) error {
	errList := validation.ValidatePod(pod)

	// remove error for no containers if requested
	if ignoreContainers {
		errList = errList.Filter(func(e error) bool {
			return e.Error() == "spec.containers: Required value"
		})
	}
	return errList.ToAggregate()
}

func deployWithPod(meta kube.ObjectMeta, attached Entity) (*deploy.Deployment, error) {
	pod, err := newDefaultPod(meta, attached.Source())
	if err != nil {
		return nil, err
	}

	err = pod.Attach(attached)
	if err != nil {
		return nil, err
	}

	return pod.Deployment()
}
<commit_msg>implemented pod Deployment()<commit_after>package entity

import (
	"fmt"

	"rsprd.com/spread/pkg/deploy"
	"rsprd.com/spread/pkg/image"

	kube "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/validation"
)

var DefaultPodSpec = kube.PodSpec{
	RestartPolicy: kube.RestartPolicyAlways,
	DNSPolicy:     kube.DNSDefault,
}

// Pod represents kube.Pod in the Redspread hierarchy.
type Pod struct {
	base
	pod        *kube.Pod
	containers []*Container
}

func NewPod(kubePod *kube.Pod, defaults kube.ObjectMeta, source string, objects ...deploy.KubeObject) (*Pod, error) {
	if kubePod == nil {
		return nil, fmt.Errorf("cannot create Pod from nil `%s`", source)
	}

	// deep copy
	kubePod, err := copyPod(kubePod)
	if err != nil {
		return nil, err
	}

	base, err := newBase(EntityPod, defaults, source, objects)
	if err != nil {
		return nil, fmt.Errorf("could not create Pod from `%s`: %v", source, err)
	}

	pod := Pod{base: base}
	for _, v := range kubePod.Spec.Containers {
		container, err := NewContainer(v, defaults, source)
		if err != nil {
			return nil, err
		} else {
			pod.containers = append(pod.containers, container)
		}
	}
	kubePod.Spec.Containers = []kube.Container{}

	base.setDefaults(kubePod)
	if err := validatePod(kubePod, true); err != nil {
		return nil, fmt.Errorf("could not create Pod from `%s`: %v", source, err)
	}

	pod.pod = kubePod
	return &pod, nil
}

func NewPodFromPodSpec(meta kube.ObjectMeta, podSpec kube.PodSpec, defaults kube.ObjectMeta, source string, objects ...deploy.KubeObject) (*Pod, error) {
	pod := kube.Pod{
		ObjectMeta: meta,
		Spec:       podSpec,
	}
	return NewPod(&pod, defaults, source, objects...)
}

func newDefaultPod(meta kube.ObjectMeta, source string) (*Pod, error) {
	return NewPodFromPodSpec(meta, DefaultPodSpec, kube.ObjectMeta{}, source)
}

func (c *Pod) Deployment() (*deploy.Deployment, error) {
	deployment := deploy.Deployment{}

	// create Pod from tree of Entities
	kubePod, childObj, err := c.data()
	if err != nil {
		return nil, err
	}

	// add Pod to deployment
	err = deployment.Add(kubePod)
	if err != nil {
		return nil, err
	}

	// add own objects
	err = deployment.AddDeployment(c.objects)
	if err != nil {
		return nil, err
	}

	// add child objects
	err = deployment.AddDeployment(childObj)
	if err != nil {
		return nil, err
	}
	return &deployment, nil
}

func (c *Pod) Images() (images []*image.Image) {
	for _, v := range c.containers {
		images = append(images, v.Images()...)
	}
	return
}

func (c *Pod) Attach(curEntity Entity) error {
	if err := c.validAttach(curEntity); err != nil {
		return err
	}
	for {
		switch e := curEntity.(type) {
		case *Image:
			container, err := newDefaultContainer(e.name(), e.Source())
			if err != nil {
				return err
			}

			err = container.Attach(e)
			curEntity = container
			break
		case *Container:
			c.containers = append(c.containers, e)
			return nil
		}
	}
}

func (c *Pod) name() string {
	return c.pod.ObjectMeta.Name
}

func (c *Pod) children() (children []Entity) {
	for _, v := range c.containers {
		children = append(children, v)
	}
	return
}

func (c *Pod) data() (pod *kube.Pod, objects deploy.Deployment, err error) {
	containers := []kube.Container{}
	for _, container := range c.containers {
		kubeContainer, cObj, err := container.data()
		if err != nil {
			return nil, objects, err
		}
		containers = append(containers, kubeContainer)
		// add containers objects
		objects.AddDeployment(cObj)
	}

	if len(containers) == 0 {
		return nil, objects, ErrorEntityNotReady
	}

	// add own objects
	objects.AddDeployment(c.objects)

	pod, err = copyPod(c.pod)
	if err != nil {
		return nil, objects, err
	}

	pod.Spec.Containers = containers
	return pod, objects, nil
}

func copyPod(pod *kube.Pod) (*kube.Pod, error) {
	copy, err := kube.Scheme.DeepCopy(pod)
	if err != nil {
		return nil, err
	}

	return copy.(*kube.Pod), nil
}

func validatePod(pod *kube.Pod, ignoreContainers bool) error {
	errList := validation.ValidatePod(pod)

	// remove error for no containers if requested
	if ignoreContainers {
		errList = errList.Filter(func(e error) bool {
			return e.Error() == "spec.containers: Required value"
		})
	}
	return errList.ToAggregate()
}

func deployWithPod(meta kube.ObjectMeta, attached Entity) (*deploy.Deployment, error) {
	pod, err := newDefaultPod(meta, attached.Source())
	if err != nil {
		return nil, err
	}

	err = pod.Attach(attached)
	if err != nil {
		return nil, err
	}

	return pod.Deployment()
}
<|endoftext|>


<commit_before>/*
Copyright 2013 Google Inc.

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

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

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

package test

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/osutil"
)

// World defines an integration test world.
//
// It's used to run the actual Camlistore binaries (camlistored,
// camput, camget, camtool, etc) together in large tests, including
// building them, finding them, and wiring them up in an isolated way.
type World struct {
	camRoot  string // typically $GOPATH[0]/src/camlistore.org
	tempDir  string
	listener net.Listener // randomly chosen 127.0.0.1 port for the server
	port     int

	server    *exec.Cmd
	isRunning int32 // state of the camlistored server. Access with sync/atomic only.
	serverErr error

	cammount *os.Process
}

// NewWorld returns a new test world.
// It requires that GOPATH is set to find the "camlistore.org" root.
func NewWorld() (*World, error) {
	if os.Getenv("GOPATH") == "" {
		return nil, errors.New("GOPATH environment variable isn't set; required to run Camlistore integration tests")
	}
	root, err := osutil.GoPackagePath("camlistore.org")
	if err == os.ErrNotExist {
		return nil, errors.New("Directory \"camlistore.org\" not found under GOPATH/src; can't run Camlistore integration tests.")
	}
	if err != nil {
		return nil, fmt.Errorf("Error searching for \"camlistore.org\" under GOPATH: %v", err)
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}

	return &World{
		camRoot:  root,
		listener: ln,
		port:     ln.Addr().(*net.TCPAddr).Port,
	}, nil
}

func (w *World) Addr() string {
	return w.listener.Addr().String()
}

// Start builds the Camlistore binaries and starts a server.
func (w *World) Start() error {
	var err error
	w.tempDir, err = ioutil.TempDir("", "camlistore-test-")
	if err != nil {
		return err
	}

	// Build.
	{
		targs := []string{
			"camget",
			"camput",
			"camtool",
			"camlistored",
		}
		// TODO(mpl): investigate why we still rebuild camlistored everytime if run through devcam test.
		// it looks like it's because we always resync the UI files and hence redo the embeds. Next CL.
		var latestModtime time.Time
		for _, target := range targs {
			binPath := filepath.Join(w.camRoot, "bin", target)
			fi, err := os.Stat(binPath)
			if err != nil {
				if !os.IsNotExist(err) {
					return fmt.Errorf("could not stat %v: %v", binPath, err)
				}
			} else {
				modTime := fi.ModTime()
				if modTime.After(latestModtime) {
					latestModtime = modTime
				}
			}
		}
		// TODO(mpl): when running with -v (either with go test or devcam test), append it for make.go as well
		cmd := exec.Command("go", "run", "make.go",
			fmt.Sprintf("--if_mods_since=%d", latestModtime.Unix()),
		)
		cmd.Dir = w.camRoot
		log.Print("Running make.go to build camlistore binaries for testing...")
		out, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("Error building world: %v, %s", err, string(out))
		}
		log.Print("Ran make.go.")
	}

	// Start camlistored.
	{
		w.server = exec.Command(
			filepath.Join(w.camRoot, "bin", "camlistored"),
			"--openbrowser=false",
			"--configfile="+filepath.Join(w.camRoot, "pkg", "test", "testdata", "server-config.json"),
			"--listen=FD:3",
			"--pollparent=true",
		)
		var buf bytes.Buffer
		w.server.Stdout = &buf
		w.server.Stderr = &buf
		w.server.Dir = w.tempDir
		w.server.Env = append(os.Environ(),
			"CAMLI_DEBUG=1",
			"CAMLI_ROOT="+w.tempDir,
			"CAMLI_SECRET_RING="+filepath.Join(w.camRoot, filepath.FromSlash("pkg/jsonsign/testdata/test-secring.gpg")),
			"CAMLI_BASE_URL=http://127.0.0.1:"+strconv.Itoa(w.port),
		)
		listenerFD, err := w.listener.(*net.TCPListener).File()
		if err != nil {
			return err
		}
		w.server.ExtraFiles = []*os.File{listenerFD}
		if err := w.server.Start(); err != nil {
			w.serverErr = fmt.Errorf("starting camlistored: %v", err)
			return w.serverErr
		}
		atomic.StoreInt32(&w.isRunning, 1)
		waitc := make(chan error, 1)
		go func() {
			err := w.server.Wait()
			w.serverErr = fmt.Errorf("%v: %s", err, buf.String())
			atomic.StoreInt32(&w.isRunning, 0)
			waitc <- w.serverErr
		}()
		upc := make(chan bool)
		timeoutc := make(chan bool)
		go func() {
			for i := 0; i < 100; i++ {
				res, err := http.Get("http://127.0.0.1:" + strconv.Itoa(w.port))
				if err == nil {
					res.Body.Close()
					upc <- true
					return
				}
				time.Sleep(50 * time.Millisecond)
			}
			w.serverErr = errors.New(buf.String())
			atomic.StoreInt32(&w.isRunning, 0)
			timeoutc <- true
		}()

		select {
		case <-waitc:
			return fmt.Errorf("server exited: %v", w.serverErr)
		case <-timeoutc:
			return fmt.Errorf("server never became reachable: %v", w.serverErr)
		case <-upc:
			if err := w.Ping(); err != nil {
				return err
			}
			// Success.
		}
	}
	return nil
}

// Ping returns an error if the world's camlistored is not running.
func (w *World) Ping() error {
	if atomic.LoadInt32(&w.isRunning) != 1 {
		return fmt.Errorf("camlistored not running: %v", w.serverErr)
	}
	return nil
}

func (w *World) Stop() {
	if w == nil {
		return
	}
	if err := w.server.Process.Kill(); err != nil {
		log.Fatalf("killed failed: %v", err)
	}

	if d := w.tempDir; d != "" {
		os.RemoveAll(d)
	}
}

func (w *World) NewPermanode(t *testing.T) blob.Ref {
	if err := w.Ping(); err != nil {
		t.Fatal(err)
	}
	out := MustRunCmd(t, w.Cmd("camput", "permanode"))
	br, ok := blob.Parse(strings.TrimSpace(out))
	if !ok {
		t.Fatalf("Expected permanode in camput stdout; got %q", out)
	}
	return br
}

func (w *World) Cmd(binary string, args ...string) *exec.Cmd {
	return w.CmdWithEnv(binary, os.Environ(), args...)
}

func (w *World) CmdWithEnv(binary string, env []string, args ...string) *exec.Cmd {
	hasVerbose := func() bool {
		for _, v := range args {
			if v == "-verbose" || v == "--verbose" {
				return true
			}
		}
		return false
	}
	var cmd *exec.Cmd
	switch binary {
	case "camget", "camput", "camtool", "cammount":
		// TODO(mpl): lift the camput restriction when we have a unified logging mechanism
		if binary == "camput" && !hasVerbose() {
			// camput and camtool are the only ones to have a -verbose flag through cmdmain
			// but camtool is never used. (and cammount does not even have a -verbose).
			args = append([]string{"-verbose"}, args...)
		}
		cmd = exec.Command(filepath.Join(w.camRoot, "bin", binary), args...)
		clientConfigDir := filepath.Join(w.camRoot, "config", "dev-client-dir")
		cmd.Env = append([]string{
			"CAMLI_CONFIG_DIR=" + clientConfigDir,
			// Respected by env expansions in config/dev-client-dir/client-config.json:
			"CAMLI_SERVER=" + w.ServerBaseURL(),
			"CAMLI_SECRET_RING=" + w.SecretRingFile(),
			"CAMLI_KEYID=" + w.ClientIdentity(),
			"CAMLI_AUTH=userpass:testuser:passTestWorld",
		}, env...)
	default:
		panic("Unknown binary " + binary)
	}
	return cmd
}

func (w *World) ServerBaseURL() string {
	return fmt.Sprintf("http://127.0.0.1:%d", w.port)
}

var theWorld *World

// GetWorld returns (creating if necessary) a test singleton world.
// It calls Fatal on the provided test if there are problems.
func GetWorld(t *testing.T) *World {
	w := theWorld
	if w == nil {
		var err error
		w, err = NewWorld()
		if err != nil {
			t.Fatalf("Error finding test world: %v", err)
		}
		err = w.Start()
		if err != nil {
			t.Fatalf("Error starting test world: %v", err)
		}
		theWorld = w
	}
	return w
}

// GetWorldMaybe returns the current World. It might be nil.
func GetWorldMaybe(t *testing.T) *World {
	return theWorld
}

// RunCmd runs c (which is assumed to be something short-lived, like a
// camput or camget command), capturing its stdout for return, and
// also capturing its stderr, just in the case of errors.
// If there's an error, the return error fully describes the command and
// all output.
func RunCmd(c *exec.Cmd) (output string, err error) {
	var stdout, stderr bytes.Buffer
	c.Stderr = &stderr
	c.Stdout = &stdout
	err = c.Run()
	if err != nil {
		return "", fmt.Errorf("Error running command %+v: Stdout:\n%s\nStderr:\n%s\n", c, stdout.String(), stderr.String())
	}
	return stdout.String(), nil
}

// MustRunCmd wraps RunCmd, failing t if RunCmd returns an error.
func MustRunCmd(t *testing.T, c *exec.Cmd) string {
	out, err := RunCmd(c)
	if err != nil {
		t.Fatal(err)
	}
	return out
}

// ClientIdentity returns the GPG identity to use in World tests, suitable
// for setting in CAMLI_KEYID.
func (w *World) ClientIdentity() string {
	return "26F5ABDA"
}

// SecretRingFile returns the GnuPG secret ring, suitable for setting
// in CAMLI_SECRET_RING.
func (w *World) SecretRingFile() string {
	return filepath.Join(w.camRoot, "pkg", "jsonsign", "testdata", "test-secring.gpg")
}

// SearchHandlerPath returns the path to the search handler, with trailing slash.
func (w *World) SearchHandlerPath() string { return "/my-search/" }
<commit_msg>test/world: print more info when testing.Verbose()<commit_after>/*
Copyright 2013 Google Inc.

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

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

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

package test

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/osutil"
)

// World defines an integration test world.
//
// It's used to run the actual Camlistore binaries (camlistored,
// camput, camget, camtool, etc) together in large tests, including
// building them, finding them, and wiring them up in an isolated way.
type World struct {
	camRoot  string // typically $GOPATH[0]/src/camlistore.org
	tempDir  string
	listener net.Listener // randomly chosen 127.0.0.1 port for the server
	port     int

	server    *exec.Cmd
	isRunning int32 // state of the camlistored server. Access with sync/atomic only.
	serverErr error

	cammount *os.Process
}

// NewWorld returns a new test world.
// It requires that GOPATH is set to find the "camlistore.org" root.
func NewWorld() (*World, error) {
	if os.Getenv("GOPATH") == "" {
		return nil, errors.New("GOPATH environment variable isn't set; required to run Camlistore integration tests")
	}
	root, err := osutil.GoPackagePath("camlistore.org")
	if err == os.ErrNotExist {
		return nil, errors.New("Directory \"camlistore.org\" not found under GOPATH/src; can't run Camlistore integration tests.")
	}
	if err != nil {
		return nil, fmt.Errorf("Error searching for \"camlistore.org\" under GOPATH: %v", err)
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}

	return &World{
		camRoot:  root,
		listener: ln,
		port:     ln.Addr().(*net.TCPAddr).Port,
	}, nil
}

func (w *World) Addr() string {
	return w.listener.Addr().String()
}

// Start builds the Camlistore binaries and starts a server.
func (w *World) Start() error {
	var err error
	w.tempDir, err = ioutil.TempDir("", "camlistore-test-")
	if err != nil {
		return err
	}

	// Build.
	{
		targs := []string{
			"camget",
			"camput",
			"camtool",
			"camlistored",
		}
		// TODO(mpl): investigate why we still rebuild camlistored everytime if run through devcam test.
		// it looks like it's because we always resync the UI files and hence redo the embeds. Next CL.
		var latestModtime time.Time
		for _, target := range targs {
			binPath := filepath.Join(w.camRoot, "bin", target)
			fi, err := os.Stat(binPath)
			if err != nil {
				if !os.IsNotExist(err) {
					return fmt.Errorf("could not stat %v: %v", binPath, err)
				}
			} else {
				modTime := fi.ModTime()
				if modTime.After(latestModtime) {
					latestModtime = modTime
				}
			}
		}
		cmd := exec.Command("go", "run", "make.go",
			fmt.Sprintf("--if_mods_since=%d", latestModtime.Unix()),
		)
		if testing.Verbose() {
			// TODO(mpl): do the same when -verbose with devcam test. Even better: see if testing.Verbose
			// can be made true if devcam test -verbose ?
			cmd.Args = append(cmd.Args, "-v=true")
		}
		cmd.Dir = w.camRoot
		log.Print("Running make.go to build camlistore binaries for testing...")
		out, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("Error building world: %v, %s", err, string(out))
		}
		if testing.Verbose() {
			log.Printf("%s\n", out)
		}
		log.Print("Ran make.go.")
	}

	// Start camlistored.
	{
		w.server = exec.Command(
			filepath.Join(w.camRoot, "bin", "camlistored"),
			"--openbrowser=false",
			"--configfile="+filepath.Join(w.camRoot, "pkg", "test", "testdata", "server-config.json"),
			"--listen=FD:3",
			"--pollparent=true",
		)
		var buf bytes.Buffer
		if testing.Verbose() {
			w.server.Stdout = os.Stdout
			w.server.Stderr = os.Stderr
		} else {
			w.server.Stdout = &buf
			w.server.Stderr = &buf
		}
		w.server.Dir = w.tempDir
		w.server.Env = append(os.Environ(),
			"CAMLI_DEBUG=1",
			"CAMLI_ROOT="+w.tempDir,
			"CAMLI_SECRET_RING="+filepath.Join(w.camRoot, filepath.FromSlash("pkg/jsonsign/testdata/test-secring.gpg")),
			"CAMLI_BASE_URL=http://127.0.0.1:"+strconv.Itoa(w.port),
		)
		listenerFD, err := w.listener.(*net.TCPListener).File()
		if err != nil {
			return err
		}
		w.server.ExtraFiles = []*os.File{listenerFD}
		if err := w.server.Start(); err != nil {
			w.serverErr = fmt.Errorf("starting camlistored: %v", err)
			return w.serverErr
		}
		atomic.StoreInt32(&w.isRunning, 1)
		waitc := make(chan error, 1)
		go func() {
			err := w.server.Wait()
			w.serverErr = fmt.Errorf("%v: %s", err, buf.String())
			atomic.StoreInt32(&w.isRunning, 0)
			waitc <- w.serverErr
		}()
		upc := make(chan bool)
		timeoutc := make(chan bool)
		go func() {
			for i := 0; i < 100; i++ {
				res, err := http.Get("http://127.0.0.1:" + strconv.Itoa(w.port))
				if err == nil {
					res.Body.Close()
					upc <- true
					return
				}
				time.Sleep(50 * time.Millisecond)
			}
			w.serverErr = errors.New(buf.String())
			atomic.StoreInt32(&w.isRunning, 0)
			timeoutc <- true
		}()

		select {
		case <-waitc:
			return fmt.Errorf("server exited: %v", w.serverErr)
		case <-timeoutc:
			return fmt.Errorf("server never became reachable: %v", w.serverErr)
		case <-upc:
			if err := w.Ping(); err != nil {
				return err
			}
			// Success.
		}
	}
	return nil
}

// Ping returns an error if the world's camlistored is not running.
func (w *World) Ping() error {
	if atomic.LoadInt32(&w.isRunning) != 1 {
		return fmt.Errorf("camlistored not running: %v", w.serverErr)
	}
	return nil
}

func (w *World) Stop() {
	if w == nil {
		return
	}
	if err := w.server.Process.Kill(); err != nil {
		log.Fatalf("killed failed: %v", err)
	}

	if d := w.tempDir; d != "" {
		os.RemoveAll(d)
	}
}

func (w *World) NewPermanode(t *testing.T) blob.Ref {
	if err := w.Ping(); err != nil {
		t.Fatal(err)
	}
	out := MustRunCmd(t, w.Cmd("camput", "permanode"))
	br, ok := blob.Parse(strings.TrimSpace(out))
	if !ok {
		t.Fatalf("Expected permanode in camput stdout; got %q", out)
	}
	return br
}

func (w *World) Cmd(binary string, args ...string) *exec.Cmd {
	return w.CmdWithEnv(binary, os.Environ(), args...)
}

func (w *World) CmdWithEnv(binary string, env []string, args ...string) *exec.Cmd {
	hasVerbose := func() bool {
		for _, v := range args {
			if v == "-verbose" || v == "--verbose" {
				return true
			}
		}
		return false
	}
	var cmd *exec.Cmd
	switch binary {
	case "camget", "camput", "camtool", "cammount":
		// TODO(mpl): lift the camput restriction when we have a unified logging mechanism
		if binary == "camput" && !hasVerbose() {
			// camput and camtool are the only ones to have a -verbose flag through cmdmain
			// but camtool is never used. (and cammount does not even have a -verbose).
			args = append([]string{"-verbose"}, args...)
		}
		cmd = exec.Command(filepath.Join(w.camRoot, "bin", binary), args...)
		clientConfigDir := filepath.Join(w.camRoot, "config", "dev-client-dir")
		cmd.Env = append([]string{
			"CAMLI_CONFIG_DIR=" + clientConfigDir,
			// Respected by env expansions in config/dev-client-dir/client-config.json:
			"CAMLI_SERVER=" + w.ServerBaseURL(),
			"CAMLI_SECRET_RING=" + w.SecretRingFile(),
			"CAMLI_KEYID=" + w.ClientIdentity(),
			"CAMLI_AUTH=userpass:testuser:passTestWorld",
		}, env...)
	default:
		panic("Unknown binary " + binary)
	}
	return cmd
}

func (w *World) ServerBaseURL() string {
	return fmt.Sprintf("http://127.0.0.1:%d", w.port)
}

var theWorld *World

// GetWorld returns (creating if necessary) a test singleton world.
// It calls Fatal on the provided test if there are problems.
func GetWorld(t *testing.T) *World {
	w := theWorld
	if w == nil {
		var err error
		w, err = NewWorld()
		if err != nil {
			t.Fatalf("Error finding test world: %v", err)
		}
		err = w.Start()
		if err != nil {
			t.Fatalf("Error starting test world: %v", err)
		}
		theWorld = w
	}
	return w
}

// GetWorldMaybe returns the current World. It might be nil.
func GetWorldMaybe(t *testing.T) *World {
	return theWorld
}

// RunCmd runs c (which is assumed to be something short-lived, like a
// camput or camget command), capturing its stdout for return, and
// also capturing its stderr, just in the case of errors.
// If there's an error, the return error fully describes the command and
// all output.
func RunCmd(c *exec.Cmd) (output string, err error) {
	var stdout, stderr bytes.Buffer
	if testing.Verbose() {
		c.Stderr = os.Stdout
		c.Stdout = os.Stderr
	} else {
		c.Stderr = &stderr
		c.Stdout = &stdout
	}
	err = c.Run()
	if err != nil {
		return "", fmt.Errorf("Error running command %+v: Stdout:\n%s\nStderr:\n%s\n", c, stdout.String(), stderr.String())
	}
	return stdout.String(), nil
}

// MustRunCmd wraps RunCmd, failing t if RunCmd returns an error.
func MustRunCmd(t *testing.T, c *exec.Cmd) string {
	out, err := RunCmd(c)
	if err != nil {
		t.Fatal(err)
	}
	return out
}

// ClientIdentity returns the GPG identity to use in World tests, suitable
// for setting in CAMLI_KEYID.
func (w *World) ClientIdentity() string {
	return "26F5ABDA"
}

// SecretRingFile returns the GnuPG secret ring, suitable for setting
// in CAMLI_SECRET_RING.
func (w *World) SecretRingFile() string {
	return filepath.Join(w.camRoot, "pkg", "jsonsign", "testdata", "test-secring.gpg")
}

// SearchHandlerPath returns the path to the search handler, with trailing slash.
func (w *World) SearchHandlerPath() string { return "/my-search/" }
<|endoftext|>


<commit_before>package units

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

type unit int64

// See: http://en.wikipedia.org/wiki/Binary_prefix
const (
	// Decimal
	KB unit = 1000
	MB      = 1000 * KB
	GB      = 1000 * MB
	TB      = 1000 * GB
	PB      = 1000 * TB

	// Binary
	KiB unit = 1024
	MiB      = 1024 * KiB
	GiB      = 1024 * MiB
	TiB      = 1024 * GiB
	PiB      = 1024 * TiB
)

type unitMap map[string]unit

var (
	decimalMap = unitMap{"k": KB, "m": MB, "g": GB, "t": TB, "p": PB}
	binaryMap  = unitMap{"k": KiB, "m": MiB, "g": GiB, "t": TiB, "p": PiB}
)

var sizeRegex *regexp.Regexp

func init() {
	sizeRegex = regexp.MustCompile("^(\\d+)([kKmMgGtTpP])?[bB]?$")
}

var unitAbbrs = [...]string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}

// HumanSize returns a human-readable approximation of a size
// using SI standard (eg. "44kB", "17MB")
func HumanSize(size int64) string {
	i := 0
	sizef := float64(size)
	for sizef >= 1000.0 {
		sizef = sizef / 1000.0
		i++
	}
	return fmt.Sprintf("%.4g %s", sizef, unitAbbrs[i])
}

// FromHumanSize returns an integer from a human-readable specification of a
// size using SI standard (eg. "44kB", "17MB")
func FromHumanSize(size string) (int64, error) {
	return parseSize(size, decimalMap)
}

// Parses a human-readable string representing an amount of RAM
// in bytes, kibibytes, mebibytes, gibibytes, or tebibytes and
// returns the number of bytes, or -1 if the string is unparseable.
// Units are case-insensitive, and the 'b' suffix is optional.
func RAMInBytes(size string) (int64, error) {
	return parseSize(size, binaryMap)
}

// Parses the human-readable size string into the amount it represents
func parseSize(sizeStr string, uMap unitMap) (int64, error) {
	matches := sizeRegex.FindStringSubmatch(sizeStr)
	if len(matches) != 3 {
		return -1, fmt.Errorf("Invalid size: '%s'", sizeStr)
	}

	size, err := strconv.ParseInt(matches[1], 10, 0)
	if err != nil {
		return -1, err
	}

	unitPrefix := strings.ToLower(matches[2])
	if mul, ok := uMap[unitPrefix]; ok {
		size *= int64(mul)
	}

	return size, nil
}
<commit_msg>pkg/units: Unit constants directly int64<commit_after>package units

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

// See: http://en.wikipedia.org/wiki/Binary_prefix
const (
	// Decimal
	KB = 1000
	MB = 1000 * KB
	GB = 1000 * MB
	TB = 1000 * GB
	PB = 1000 * TB

	// Binary
	KiB = 1024
	MiB = 1024 * KiB
	GiB = 1024 * MiB
	TiB = 1024 * GiB
	PiB = 1024 * TiB
)

type unitMap map[string]int64

var (
	decimalMap = unitMap{"k": KB, "m": MB, "g": GB, "t": TB, "p": PB}
	binaryMap  = unitMap{"k": KiB, "m": MiB, "g": GiB, "t": TiB, "p": PiB}
)

var sizeRegex *regexp.Regexp

func init() {
	sizeRegex = regexp.MustCompile("^(\\d+)([kKmMgGtTpP])?[bB]?$")
}

var unitAbbrs = [...]string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}

// HumanSize returns a human-readable approximation of a size
// using SI standard (eg. "44kB", "17MB")
func HumanSize(size int64) string {
	i := 0
	sizef := float64(size)
	for sizef >= 1000.0 {
		sizef = sizef / 1000.0
		i++
	}
	return fmt.Sprintf("%.4g %s", sizef, unitAbbrs[i])
}

// FromHumanSize returns an integer from a human-readable specification of a
// size using SI standard (eg. "44kB", "17MB")
func FromHumanSize(size string) (int64, error) {
	return parseSize(size, decimalMap)
}

// Parses a human-readable string representing an amount of RAM
// in bytes, kibibytes, mebibytes, gibibytes, or tebibytes and
// returns the number of bytes, or -1 if the string is unparseable.
// Units are case-insensitive, and the 'b' suffix is optional.
func RAMInBytes(size string) (int64, error) {
	return parseSize(size, binaryMap)
}

// Parses the human-readable size string into the amount it represents
func parseSize(sizeStr string, uMap unitMap) (int64, error) {
	matches := sizeRegex.FindStringSubmatch(sizeStr)
	if len(matches) != 3 {
		return -1, fmt.Errorf("Invalid size: '%s'", sizeStr)
	}

	size, err := strconv.ParseInt(matches[1], 10, 0)
	if err != nil {
		return -1, err
	}

	unitPrefix := strings.ToLower(matches[2])
	if mul, ok := uMap[unitPrefix]; ok {
		size *= mul
	}

	return size, nil
}
<|endoftext|>


<commit_before>package main

// Copyright 2017 Microsoft Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/marstr/collection"
)

// LatestStrategy evaluates the Azure-SDK-for-Go repository for the latest available API versions of each service.
type LatestStrategy struct {
	Root      string
	Predicate func(packageName string) bool
}

// AcceptAll is a predefined value for `LatestStrategy.Predicate` which always returns true.
func AcceptAll(name string) (result bool) {
	result = true
	return
}

// IgnorePreview searches a packages "API Version" to see if it contains the word "preview". It only returns true when a package is not a preview.
func IgnorePreview(name string) (result bool) {
	matches := packageName.FindStringSubmatch(name)
	if len(matches) >= 4 {
		result = !strings.Contains(matches[3], "preview")
	}
	return
}

var packageName = regexp.MustCompile(`service[/\\](?P<provider>[\w\-\.\d]+)[/\\](?P<type>[\w\-\.\d]+)(?:[/\\]management)?[/\\](?P<version>[\d\-\w\.]+)[/\\](?P<group>[\w\d\-\.]+)`)

// Enumerate scans through the known Azure SDK for Go packages and finds each
func (latest LatestStrategy) Enumerate(cancel <-chan struct{}) collection.Enumerator {
	results := make(chan interface{})

	go func() {
		defer close(results)

		type operationGroup struct {
			provider     string
			resourceType string
			group        string
		}

		type operInfo struct {
			version string
			rawpath string
		}

		maxFound := make(map[operationGroup]operInfo)

		filepath.Walk(latest.Root, func(currentPath string, info os.FileInfo, openErr error) (err error) {
			pathMatches := packageName.FindStringSubmatch(currentPath)

			if latest.Predicate == nil {
				latest.Predicate = AcceptAll
			}

			if len(pathMatches) == 0 || !info.IsDir() || !latest.Predicate(currentPath) {
				return
			}

			version := pathMatches[3]
			currentGroup := operationGroup{
				provider:     pathMatches[1],
				resourceType: pathMatches[2],
				group:        pathMatches[4],
			}

			prev, ok := maxFound[currentGroup]
			if !ok {
				maxFound[currentGroup] = operInfo{version, currentPath}
				return
			}

			if le, _ := versionle(prev.version, version); le {
				maxFound[currentGroup] = operInfo{version, currentPath}
			}

			return
		})

		for _, entry := range maxFound {
			absolute, err := filepath.Abs(entry.rawpath)
			if err != nil {
				continue
			}

			select {
			case results <- absolute:
				// Intionally Left Blank
			case <-cancel:
				return
			}
		}
	}()

	return results
}

// ErrNotVersionString is instantiated when a string not conforming to known Azure API Version patterns is parsed is if it did.
type ErrNotVersionString string

func (err ErrNotVersionString) Error() string {
	return fmt.Sprintf("`%s` is not a recognized Azure version string", string(err))
}

var versionle = func() func(string, string) (bool, error) {
	versionPattern := regexp.MustCompile(`^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})(?:[\.\-](?P<tag>.+))?$`)

	return func(left, right string) (result bool, err error) {
		leftMatch := versionPattern.FindStringSubmatch(left)
		rightMatch := versionPattern.FindStringSubmatch(right)

		if len(leftMatch) < 3 {
			err = ErrNotVersionString(left)
			return
		}

		if len(rightMatch) < 3 {
			err = ErrNotVersionString(right)
			return
		}

		for i := 1; i <= 3; i++ {
			if leftMatch[i] == rightMatch[i] {
				continue
			}

			var leftNum, rightNum int
			leftNum, err = strconv.Atoi(leftMatch[i])
			if err != nil {
				return
			}

			rightNum, err = strconv.Atoi(rightMatch[i])
			if err != nil {
				return
			}

			if leftNum < rightNum {
				result = true
				return
			}

			result = false // This line only here for readability
			return
		}

		if len(leftMatch) == 5 && len(rightMatch) == 5 {
			result = leftMatch[4] <= rightMatch[4]
			return
		} else if len(leftMatch) == 4 {
			result = false
			return
		}
		result = true
		return
	}
}()
<commit_msg>Fixing up `versionle`<commit_after>package main

// Copyright 2017 Microsoft Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/marstr/collection"
)

// LatestStrategy evaluates the Azure-SDK-for-Go repository for the latest available API versions of each service.
type LatestStrategy struct {
	Root      string
	Predicate func(packageName string) bool
}

// AcceptAll is a predefined value for `LatestStrategy.Predicate` which always returns true.
func AcceptAll(name string) (result bool) {
	result = true
	return
}

// IgnorePreview searches a packages "API Version" to see if it contains the word "preview". It only returns true when a package is not a preview.
func IgnorePreview(name string) (result bool) {
	matches := packageName.FindStringSubmatch(name)
	if len(matches) >= 4 {
		result = !strings.Contains(matches[3], "preview")
	}
	return
}

var packageName = regexp.MustCompile(`service[/\\](?P<provider>[\w\-\.\d]+)[/\\](?P<type>[\w\-\.\d]+)(?:[/\\]management)?[/\\](?P<version>[\d\-\w\.]+)[/\\](?P<group>[\w\d\-\.]+)`)

// Enumerate scans through the known Azure SDK for Go packages and finds each
func (latest LatestStrategy) Enumerate(cancel <-chan struct{}) collection.Enumerator {
	results := make(chan interface{})

	go func() {
		defer close(results)

		type operationGroup struct {
			provider     string
			resourceType string
			group        string
		}

		type operInfo struct {
			version string
			rawpath string
		}

		maxFound := make(map[operationGroup]operInfo)

		filepath.Walk(latest.Root, func(currentPath string, info os.FileInfo, openErr error) (err error) {
			pathMatches := packageName.FindStringSubmatch(currentPath)

			if latest.Predicate == nil {
				latest.Predicate = AcceptAll
			}

			if len(pathMatches) == 0 || !info.IsDir() || !latest.Predicate(currentPath) {
				return
			}

			version := pathMatches[3]
			currentGroup := operationGroup{
				provider:     pathMatches[1],
				resourceType: pathMatches[2],
				group:        pathMatches[4],
			}

			prev, ok := maxFound[currentGroup]
			if !ok {
				maxFound[currentGroup] = operInfo{version, currentPath}
				return
			}

			if le, _ := versionle(prev.version, version); le {
				maxFound[currentGroup] = operInfo{version, currentPath}
			}

			return
		})

		for _, entry := range maxFound {
			absolute, err := filepath.Abs(entry.rawpath)
			if err != nil {
				continue
			}

			select {
			case results <- absolute:
				// Intionally Left Blank
			case <-cancel:
				return
			}
		}
	}()

	return results
}

// ErrNotVersionString is instantiated when a string not conforming to known Azure API Version patterns is parsed is if it did.
type ErrNotVersionString string

func (err ErrNotVersionString) Error() string {
	return fmt.Sprintf("`%s` is not a recognized Azure version string", string(err))
}

var versionle = func() func(string, string) (bool, error) {
	versionPattern := regexp.MustCompile(`^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})(?:[\.\-](?P<tag>.+))?$`)

	return func(left, right string) (result bool, err error) {
		leftMatch := versionPattern.FindStringSubmatch(left)
		rightMatch := versionPattern.FindStringSubmatch(right)

		if len(leftMatch) < 3 {
			err = ErrNotVersionString(left)
			return
		}

		if len(rightMatch) < 3 {
			err = ErrNotVersionString(right)
			return
		}

		for i := 1; i <= 3; i++ {
			if leftMatch[i] == rightMatch[i] {
				continue
			}

			var leftNum, rightNum int
			leftNum, err = strconv.Atoi(leftMatch[i])
			if err != nil {
				return
			}

			rightNum, err = strconv.Atoi(rightMatch[i])
			if err != nil {
				return
			}

			if leftNum < rightNum {
				result = true
				return
			}

			result = false // This line only here for readability
			return
		}

		if leftTag, rightTag := leftMatch[4], rightMatch[4]; leftTag == "" && rightTag != "" {
			result = false
		} else if leftTag != "" && rightTag != "" {
			result = leftTag <= rightTag
		} else {
			result = true
		}
		return
	}
}()
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package store

import (
	"encoding/json"
	"sync/atomic"
)

const (
	SetSuccess = iota
	SetFail
	DeleteSuccess
	DeleteFail
	CreateSuccess
	CreateFail
	UpdateSuccess
	UpdateFail
	CompareAndSwapSuccess
	CompareAndSwapFail
	GetSuccess
	GetFail
	ExpireCount
	CompareAndDeleteSuccess
	CompareAndDeleteFail
)

type Stats struct {

	// Number of get requests
	GetSuccess uint64 `json:"getsSuccess"`
	GetFail    uint64 `json:"getsFail"`

	// Number of sets requests
	SetSuccess uint64 `json:"setsSuccess"`
	SetFail    uint64 `json:"setsFail"`

	// Number of delete requests
	DeleteSuccess uint64 `json:"deleteSuccess"`
	DeleteFail    uint64 `json:"deleteFail"`

	// Number of update requests
	UpdateSuccess uint64 `json:"updateSuccess"`
	UpdateFail    uint64 `json:"updateFail"`

	// Number of create requests
	CreateSuccess uint64 `json:"createSuccess"`
	CreateFail    uint64 `json:"createFail"`

	// Number of testAndSet requests
	CompareAndSwapSuccess uint64 `json:"compareAndSwapSuccess"`
	CompareAndSwapFail    uint64 `json:"compareAndSwapFail"`

	// Number of compareAndDelete requests
	CompareAndDeleteSuccess uint64 `json:"compareAndDeleteSuccess"`
	CompareAndDeleteFail    uint64 `json:"compareAndDeleteFail"`

	ExpireCount uint64 `json:"expireCount"`

	Watchers uint64 `json:"watchers"`
}

func newStats() *Stats {
	s := new(Stats)
	return s
}

func (s *Stats) clone() *Stats {
	return &Stats{s.GetSuccess, s.GetFail, s.SetSuccess, s.SetFail,
		s.DeleteSuccess, s.DeleteFail, s.UpdateSuccess, s.UpdateFail, s.CreateSuccess,
		s.CreateFail, s.CompareAndSwapSuccess, s.CompareAndSwapFail,
		s.CompareAndDeleteSuccess, s.CompareAndDeleteFail, s.Watchers, s.ExpireCount}
}

func (s *Stats) toJson() []byte {
	b, _ := json.Marshal(s)
	return b
}

func (s *Stats) Inc(field int) {
	switch field {
	case SetSuccess:
		atomic.AddUint64(&s.SetSuccess, 1)
	case SetFail:
		atomic.AddUint64(&s.SetFail, 1)
	case CreateSuccess:
		atomic.AddUint64(&s.CreateSuccess, 1)
	case CreateFail:
		atomic.AddUint64(&s.CreateFail, 1)
	case DeleteSuccess:
		atomic.AddUint64(&s.DeleteSuccess, 1)
	case DeleteFail:
		atomic.AddUint64(&s.DeleteFail, 1)
	case GetSuccess:
		atomic.AddUint64(&s.GetSuccess, 1)
	case GetFail:
		atomic.AddUint64(&s.GetFail, 1)
	case UpdateSuccess:
		atomic.AddUint64(&s.UpdateSuccess, 1)
	case UpdateFail:
		atomic.AddUint64(&s.UpdateFail, 1)
	case CompareAndSwapSuccess:
		atomic.AddUint64(&s.CompareAndSwapSuccess, 1)
	case CompareAndSwapFail:
		atomic.AddUint64(&s.CompareAndSwapFail, 1)
	case CompareAndDeleteSuccess:
		atomic.AddUint64(&s.CompareAndDeleteSuccess, 1)
	case CompareAndDeleteFail:
		atomic.AddUint64(&s.CompareAndDeleteFail, 1)
	case ExpireCount:
		atomic.AddUint64(&s.ExpireCount, 1)
	}
}
<commit_msg>store: fixed clone error for store stats.<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package store

import (
	"encoding/json"
	"sync/atomic"
)

const (
	SetSuccess = iota
	SetFail
	DeleteSuccess
	DeleteFail
	CreateSuccess
	CreateFail
	UpdateSuccess
	UpdateFail
	CompareAndSwapSuccess
	CompareAndSwapFail
	GetSuccess
	GetFail
	ExpireCount
	CompareAndDeleteSuccess
	CompareAndDeleteFail
)

type Stats struct {

	// Number of get requests
	GetSuccess uint64 `json:"getsSuccess"`
	GetFail    uint64 `json:"getsFail"`

	// Number of sets requests
	SetSuccess uint64 `json:"setsSuccess"`
	SetFail    uint64 `json:"setsFail"`

	// Number of delete requests
	DeleteSuccess uint64 `json:"deleteSuccess"`
	DeleteFail    uint64 `json:"deleteFail"`

	// Number of update requests
	UpdateSuccess uint64 `json:"updateSuccess"`
	UpdateFail    uint64 `json:"updateFail"`

	// Number of create requests
	CreateSuccess uint64 `json:"createSuccess"`
	CreateFail    uint64 `json:"createFail"`

	// Number of testAndSet requests
	CompareAndSwapSuccess uint64 `json:"compareAndSwapSuccess"`
	CompareAndSwapFail    uint64 `json:"compareAndSwapFail"`

	// Number of compareAndDelete requests
	CompareAndDeleteSuccess uint64 `json:"compareAndDeleteSuccess"`
	CompareAndDeleteFail    uint64 `json:"compareAndDeleteFail"`

	ExpireCount uint64 `json:"expireCount"`

	Watchers uint64 `json:"watchers"`
}

func newStats() *Stats {
	s := new(Stats)
	return s
}

func (s *Stats) clone() *Stats {
	return &Stats{
		GetSuccess:              s.GetSuccess,
		GetFail:                 s.GetFail,
		SetSuccess:              s.SetSuccess,
		SetFail:                 s.SetFail,
		DeleteSuccess:           s.DeleteSuccess,
		DeleteFail:              s.DeleteFail,
		UpdateSuccess:           s.UpdateSuccess,
		UpdateFail:              s.UpdateFail,
		CreateSuccess:           s.CreateSuccess,
		CreateFail:              s.CreateFail,
		CompareAndSwapSuccess:   s.CompareAndSwapSuccess,
		CompareAndSwapFail:      s.CompareAndSwapFail,
		CompareAndDeleteSuccess: s.CompareAndDeleteSuccess,
		CompareAndDeleteFail:    s.CompareAndDeleteFail,
		ExpireCount:             s.ExpireCount,
		Watchers:                s.Watchers,
	}
}

func (s *Stats) toJson() []byte {
	b, _ := json.Marshal(s)
	return b
}

func (s *Stats) Inc(field int) {
	switch field {
	case SetSuccess:
		atomic.AddUint64(&s.SetSuccess, 1)
	case SetFail:
		atomic.AddUint64(&s.SetFail, 1)
	case CreateSuccess:
		atomic.AddUint64(&s.CreateSuccess, 1)
	case CreateFail:
		atomic.AddUint64(&s.CreateFail, 1)
	case DeleteSuccess:
		atomic.AddUint64(&s.DeleteSuccess, 1)
	case DeleteFail:
		atomic.AddUint64(&s.DeleteFail, 1)
	case GetSuccess:
		atomic.AddUint64(&s.GetSuccess, 1)
	case GetFail:
		atomic.AddUint64(&s.GetFail, 1)
	case UpdateSuccess:
		atomic.AddUint64(&s.UpdateSuccess, 1)
	case UpdateFail:
		atomic.AddUint64(&s.UpdateFail, 1)
	case CompareAndSwapSuccess:
		atomic.AddUint64(&s.CompareAndSwapSuccess, 1)
	case CompareAndSwapFail:
		atomic.AddUint64(&s.CompareAndSwapFail, 1)
	case CompareAndDeleteSuccess:
		atomic.AddUint64(&s.CompareAndDeleteSuccess, 1)
	case CompareAndDeleteFail:
		atomic.AddUint64(&s.CompareAndDeleteFail, 1)
	case ExpireCount:
		atomic.AddUint64(&s.ExpireCount, 1)
	}
}
<|endoftext|>


<commit_before>// Package store provides a simple distributed key-value store. The keys and
// associated values are changed via distributed consensus, meaning that the
// values are changed only when a majority of nodes in the cluster agree on
// the new value.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
)

type command struct {
	Op    string `json:"op,omitempty"`
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

// Store is a simple key-value store, where all changes are made via Raft consensus.
type Store struct {
	RaftDir  string
	RaftBind string

	mu sync.Mutex
	m  map[string]string // The key-value store for the system.

	raft *raft.Raft // The consensus mechanism

	logger *log.Logger
}

// New returns a new Store.
func New() *Store {
	return &Store{
		m:      make(map[string]string),
		logger: log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, then this node become the first node,
// and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Setup Raft configuration.
	config := raft.DefaultConfig()

	if enableSingle {
		// Experiment with single-node to start.
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	addr, err := net.ResolveTCPAddr("tcp", s.RaftBind)
	if err != nil {
		return err
	}
	transport, err := raft.NewTCPTransport(s.RaftBind, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	// Create peer storage.
	peerStore := raft.NewJSONPeers(filepath.Join(s.RaftDir, "peers.json"), transport)

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.RaftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Create the snapshot store.
	snapshots, err := raft.NewFileSnapshotStore(s.RaftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create raft log.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra
	return nil
}

// Get returns the value for the given key.
func (s *Store) Get(key string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.m[key], nil
}

// Set sets the value for the given key.
func (s *Store) Set(key, value string) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

	c := &command{
		Op:    "set",
		Key:   key,
		Value: value,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	if err, ok := f.(error); ok {
		return err
	}

	return nil
}

// Delete deletes the given key.
func (s *Store) Delete(key string) error {
	return nil
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received join request for remote node as %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

// Apply applies a Raft log entry to the key-value store.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	switch c.Op {
	case "set":
		return f.applySet(c.Key, c.Value)
	case "delete":
		return f.applyDelete(c.Key)
	default:
		panic(fmt.Sprintf("unrecognized command op: %s", c.Op))
	}
}

// Snapshot returns a snapshot of the key-value store.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	return nil, nil
}

// Restore stores the key-value store to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	return nil
}

func (f *fsm) applySet(key, value string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.m[key] = value
	return nil
}

func (f *fsm) applyDelete(key string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	delete(f.m, key)
	return nil
}
<commit_msg>peers file path set correctly and read on start-up<commit_after>// Package store provides a simple distributed key-value store. The keys and
// associated values are changed via distributed consensus, meaning that the
// values are changed only when a majority of nodes in the cluster agree on
// the new value.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
)

type command struct {
	Op    string `json:"op,omitempty"`
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

// Store is a simple key-value store, where all changes are made via Raft consensus.
type Store struct {
	RaftDir  string
	RaftBind string

	mu sync.Mutex
	m  map[string]string // The key-value store for the system.

	raft *raft.Raft // The consensus mechanism

	logger *log.Logger
}

// New returns a new Store.
func New() *Store {
	return &Store{
		m:      make(map[string]string),
		logger: log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomesthe first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Check for any existing peers.
	peers, err := readPeersJSON(filepath.Join(s.RaftDir, "peers.json"))
	if err != nil {
		return err
	}

	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	addr, err := net.ResolveTCPAddr("tcp", s.RaftBind)
	if err != nil {
		return err
	}
	transport, err := raft.NewTCPTransport(s.RaftBind, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.RaftDir, transport)

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.RaftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Create the snapshot store.
	snapshots, err := raft.NewFileSnapshotStore(s.RaftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create raft log.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra
	return nil
}

// Get returns the value for the given key.
func (s *Store) Get(key string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.m[key], nil
}

// Set sets the value for the given key.
func (s *Store) Set(key, value string) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

	c := &command{
		Op:    "set",
		Key:   key,
		Value: value,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	if err, ok := f.(error); ok {
		return err
	}

	return nil
}

// Delete deletes the given key.
func (s *Store) Delete(key string) error {
	return nil
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received join request for remote node as %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

// Apply applies a Raft log entry to the key-value store.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	switch c.Op {
	case "set":
		return f.applySet(c.Key, c.Value)
	case "delete":
		return f.applyDelete(c.Key)
	default:
		panic(fmt.Sprintf("unrecognized command op: %s", c.Op))
	}
}

// Snapshot returns a snapshot of the key-value store.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	return nil, nil
}

// Restore stores the key-value store to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	return nil
}

func (f *fsm) applySet(key, value string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.m[key] = value
	return nil
}

func (f *fsm) applyDelete(key string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	delete(f.m, key)
	return nil
}

func readPeersJSON(path string) ([]string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	if len(b) == 0 {
		return nil, nil
	}

	var peers []string
	dec := json.NewDecoder(bytes.NewReader(b))
	if err := dec.Decode(&peers); err != nil {
		return nil, err
	}

	return peers, nil
}
<|endoftext|>


<commit_before>package store

import (
	"fmt"
	"github.com/bradylove/mescal/msg"
	"io"
	"sync"
)

type object struct {
	key    string
	value  string
	expiry int64
}

type objects map[string]object

type message struct {
	*msg.Command

	writer io.Writer
}

type Store struct {
	objects objects
	wg      sync.WaitGroup
	msgs    chan message
}

const (
	workerCount = 20
)

func NewStore() *Store {
	s := Store{
		objects: objects{
			"foo": {"foo", "bar", int64(1234)},
			"fuz": {"fuz", "baz", int64(1234)},
		},
		msgs: make(chan message),
	}

	s.setupPool()

	return &s
}

func (s *Store) setupPool() {
	s.wg.Add(workerCount)

	for i := 0; i < workerCount; i++ {
		go s.manageStore()
	}
}

func (s *Store) manageStore() {
StoreHandler:
	for {
		m, ok := <-s.msgs
		if !ok {
			break StoreHandler
		}

		switch sb := m.SubCommand.(type) {
		case msg.GetCommand:
			s.handleGetCommand(m.Command, sb, m.writer)
		default:
			fmt.Println("Unknown command")
		}
	}

	s.wg.Done()
}

func (s *Store) handleGetCommand(cmd *msg.Command, sb msg.GetCommand, w io.Writer) {
	obj := s.objects[sb.Key]

	res := msg.NewResult(
		cmd.Id,
		msg.StatusSuccess,
		msg.NewGetResult(obj.key, obj.value, obj.expiry),
	)

	res.Encode(w)
}

func (s *Store) HandleCommand(cmd *msg.Command, w io.Writer) {
	s.msgs <- message{cmd, w}
}

func (s *Store) Wait() {
	s.wg.Wait()
}

func (s *Store) Close() {
	close(s.msgs)
}
<commit_msg>Adding the Set command to the store<commit_after>package store

import (
	"fmt"
	"github.com/bradylove/mescal/msg"
	"io"
	"sync"
)

type object struct {
	key    string
	value  string
	expiry int64
}

type objects map[string]object

type message struct {
	*msg.Command

	writer io.Writer
}

type Store struct {
	objects objects
	wg      sync.WaitGroup
	msgs    chan message
}

const (
	workerCount = 20
)

func NewStore() *Store {
	s := Store{
		objects: objects{
			"foo": {"foo", "bar", int64(1234)},
			"fuz": {"fuz", "baz", int64(1234)},
		},
		msgs: make(chan message),
	}

	s.setupPool()

	return &s
}

func (s *Store) setupPool() {
	s.wg.Add(workerCount)

	for i := 0; i < workerCount; i++ {
		go s.manageStore()
	}
}

func (s *Store) manageStore() {
StoreHandler:
	for {
		m, ok := <-s.msgs
		if !ok {
			break StoreHandler
		}
		// TODO do we need to add a case statement for the 'set' cmd
		// case msg.SetCommand:
		//   s.handleSetCommand(m.Command, sb, m.Writer)
		switch sb := m.SubCommand.(type) {
		case msg.GetCommand:
			s.handleGetCommand(m.Command, sb, m.writer) // Should the m.writer be m.Writer?
		default:
			fmt.Println("Unknown command")
		}
	}

	s.wg.Done()
}

func (s *Store) handleGetCommand(cmd *msg.Command, sb msg.GetCommand, w io.Writer) {
	obj := s.objects[sb.Key]

	res := msg.NewResult(
		cmd.Id,
		msg.StatusSuccess,
		msg.NewGetResult(obj.key, obj.value, obj.expiry),
	)

	res.Encode(w)
}

func (s *Store) handleSetCommand(cmd *msg.Command, sb msg.SetCommand, w io.Writer) {
	obj := s.objects[sb.Key]

	res := msg.NewResult(
		cmd.Id,
		msg.StatusSuccess,
		msg.NewSetResult(obj.key, obj.value, obj.expiry),
	)

	res.Encode(w)
}
func (s *Store) HandleCommand(cmd *msg.Command, w io.Writer) {
	s.msgs <- message{cmd, w}
}

func (s *Store) Wait() {
	s.wg.Wait()
}

func (s *Store) Close() {
	close(s.msgs)
}
<|endoftext|>