<commit_before>package eval

import (
	"fmt"
	"jacob/black/ast"
	"jacob/black/object"
	"jacob/black/token"
	"math"
)

var (
	// ConstTrue is the only and only true value
	ConstTrue = &object.Boolean{Value: true}
	// ConstFalse is the only and only false value
	ConstFalse = &object.Boolean{Value: false}
	// ConstNil is the only and only nil
	ConstNil = &object.Nil{}
)

func newError(pos token.Position, format string, v ...interface{}) *object.Error {
	return &object.Error{Message: fmt.Sprintf(format, v...), Pos: pos}
}

func isError(o object.Object) bool {
	if o != nil {
		return o.Type() == object.ErrorType
	}
	return false
}

// Eval evaluates the program node and returns an object as a result
func Eval(node ast.Node, env *object.Environment) object.Object {
	switch node := node.(type) {

	// statements
	case *ast.Program:
		return evalProgram(node, env)
	case *ast.LetStatement:
		val := Eval(node.Value, env)
		if isError(val) {
			return val
		}
		env.Set(node.Name.Value, val)
	case *ast.ReturnStatement:
		val := Eval(node.Value, env)
		if isError(val) {
			return val
		}
		return &object.ReturnValue{Value: val}
	case *ast.BlockStatement:
		return evalBlockStatement(node, env)
	case *ast.ExpressionStatement:
		return Eval(node.Expression, env)

		// expressions
	case *ast.PrefixExpression:
		right := Eval(node.Right, env)
		if isError(right) {
			return right
		}
		return evalPrefixExpr(node.Token, right)
	case *ast.InfixExpression:
		left := Eval(node.Left, env)
		if isError(left) {
			return left
		}
		right := Eval(node.Right, env)
		if isError(right) {
			return right
		}
		return evalInfixExpr(node.Token, left, right)
	case *ast.IfExpression:
		return evalIfExpr(node, env)

		// literals
	case *ast.IntegerLiteral:
		return &object.Integer{Value: node.Value}
	case *ast.FloatLiteral:
		return &object.Float{Value: node.Value}
	case *ast.BooleanLiteral:
		return boolToBoolean(node.Value)

	case *ast.Identifier:
		return evalIdentifier(node, env)
	}
	return nil
}

func evalProgram(program *ast.Program, env *object.Environment) object.Object {
	var result object.Object

	for _, s := range program.Statements {
		result = Eval(s, env)

		if result != nil {
			// pass up the return type to the top level
			switch result.Type() {
			case object.ReturnType:
				// unwrap the return value
				return result.(*object.ReturnValue).Value
			case object.ErrorType:
				return result.(*object.Error)
			}
		}
	}

	return result
}

func evalBlockStatement(block *ast.BlockStatement, env *object.Environment) object.Object {
	var result object.Object

	for _, s := range block.Statements {
		result = Eval(s, env)

		if result != nil {
			if result.Type() == object.ReturnType || result.Type() == object.ErrorType {
				return result
			}
		}
	}

	return result
}

func evalIfExpr(node *ast.IfExpression, env *object.Environment) object.Object {
	cond := Eval(node.Cond, env)

	if isError(cond) {
		return cond
	}

	if isTruthy(cond) {
		return Eval(node.Do, env)
	} else if node.Else != nil {
		return Eval(node.Else, env)
	}

	return ConstNil
}

// isTruthy - everything is true execpt for false and nil
func isTruthy(o object.Object) bool {
	switch o {
	case ConstFalse, ConstNil:
		return false
	default:
		// special case: 0 or 0.0 is not truthy
		switch o.Type() {
		case object.IntType:
			if o.(*object.Integer).Value == 0 {
				return false
			}
		case object.FloatType:
			if o.(*object.Float).Value == 0.0 {
				return false
			}
		}

		return true
	}
}

func boolToBoolean(b bool) *object.Boolean {
	if b {
		return ConstTrue
	}
	return ConstFalse
}

func evalPrefixExpr(op token.Token, right object.Object) object.Object {
	switch op.Type {
	case token.Bang:
		return evalBangOperatorExpr(right)
	case token.Minus:
		return evalMinusPrefixOperatorExpr(op.Pos, right)
	default:
		return newError(op.Pos, "unknown operator '%s' for type '%s'", op, right.Type())
	}
}

func evalInfixExpr(op token.Token, left object.Object, right object.Object) object.Object {

	if !left.CanApply(op.Type, right.Type()) {
		return newError(op.Pos, "cannot apply operator '%s' for type '%s' and '%s'", op, left.Type(), right.Type())
	}

	// test and convert int to float if needed
	if left.Type() == object.IntType && right.Type() == object.IntType {
		// both ints. easy
		return evalIntegerInfixExpr(op, left, right)
	} else if left.Type() == object.FloatType && right.Type() == object.FloatType {
		// both floats. easy
		return evalFloatInfixExpr(op, left, right)
	}

	// one of them must be a int and the other a float
	if left.Type() == object.IntType || right.Type() == object.IntType {
		if left.Type() == object.FloatType {
			// left is float, right is int.
			// premote right to float
			val := right.(*object.Integer).Value
			right = &object.Float{Value: float64(val)}
			return evalFloatInfixExpr(op, left, right)

		} else if right.Type() == object.FloatType {
			// right is float, left is int
			// premote left to float
			val := left.(*object.Integer).Value
			left = &object.Float{Value: float64(val)}

			return evalFloatInfixExpr(op, left, right)
		}
	}

	// compare actual runtime object
	if op.Type == token.Equal {
		return boolToBoolean(left == right)
	} else if op.Type == token.NotEqual {
		return boolToBoolean(left != right)
	}

	// otherwise 2 objects that don't match
	return newError(op.Pos, "unknown operator '%s' for type '%s' and '%s'", op, left.Type(), right.Type())
}

func evalIntegerInfixExpr(op token.Token, left object.Object, right object.Object) object.Object {
	leftVal := left.(*object.Integer).Value
	rightVal := right.(*object.Integer).Value

	switch op.Type {
	case token.Plus:
		return &object.Integer{Value: leftVal + rightVal}
	case token.Minus:
		return &object.Integer{Value: leftVal - rightVal}
	case token.Times:
		return &object.Integer{Value: leftVal * rightVal}
	case token.Divide:
		if rightVal == 0 {
			return newError(op.Pos, "cannot divide %d by 0", leftVal)
		}
		return &object.Integer{Value: leftVal / rightVal}
	case token.Exp:
		return &object.Integer{Value: int64(math.Pow(float64(leftVal), float64(rightVal)))}
	case token.Mod:
		if rightVal == 0 {
			return newError(op.Pos, "cannot modulo %d by 0", leftVal)
		}
		return &object.Integer{Value: leftVal % rightVal}
	case token.Less:
		return boolToBoolean(leftVal < rightVal)
	case token.Greater:
		return boolToBoolean(leftVal > rightVal)
	case token.Equal:
		return boolToBoolean(leftVal == rightVal)
	case token.NotEqual:
		return boolToBoolean(leftVal != rightVal)
	default:
		return newError(op.Pos, "unknown operator '%s' for type '%s' and '%s'", op.Type, left.Type(), right.Type())
	}
}

func evalFloatInfixExpr(op token.Token, left object.Object, right object.Object) object.Object {
	leftVal := left.(*object.Float).Value
	rightVal := right.(*object.Float).Value

	switch op.Type {
	case token.Plus:
		return &object.Float{Value: leftVal + rightVal}
	case token.Minus:
		return &object.Float{Value: leftVal - rightVal}
	case token.Times:
		return &object.Float{Value: leftVal * rightVal}
	case token.Divide:
		if rightVal == 0 {
			return newError(op.Pos, "cannot divide %d by 0", leftVal)
		}
		return &object.Float{Value: leftVal / rightVal}
	case token.Exp:
		return &object.Float{Value: math.Pow(leftVal, rightVal)}
	case token.Mod:
		if rightVal == 0 {
			return newError(op.Pos, "cannot modulo %d by 0", leftVal)
		}
		return &object.Float{Value: math.Mod(leftVal, rightVal)}
	case token.Less:
		return boolToBoolean(leftVal < rightVal)
	case token.Greater:
		return boolToBoolean(leftVal > rightVal)
	case token.Equal:
		return boolToBoolean(leftVal == rightVal)
	case token.NotEqual:
		return boolToBoolean(leftVal != rightVal)
	default:
		return newError(op.Pos, "unknown operator '%s' for type '%s' and '%s'", op.Type, left.Type(), right.Type())
	}
}

func evalBangOperatorExpr(right object.Object) object.Object {
	switch right {
	case ConstTrue:
		return ConstFalse
	case ConstFalse:
		return ConstTrue
	case ConstNil:
		return ConstTrue
	default:
		return boolToBoolean(!isTruthy(right))
	}
}

func evalMinusPrefixOperatorExpr(pos token.Position, right object.Object) object.Object {
	switch right.Type() {
	case object.IntType:
		v := right.(*object.Integer).Value
		return &object.Integer{Value: -v}
	case object.FloatType:
		v := right.(*object.Float).Value
		return &object.Float{Value: -v}
	default:
		return newError(pos, "unknown operator '-' for type '%s'", right.Type())
	}
}

func evalIdentifier(id *ast.Identifier, env *object.Environment) object.Object {
	if val, ok := env.Get(id.Value); ok {
		return val
	}
	return newError(id.Token.Pos, "identifier not found: %s", id.Value)
}
<commit_msg>added = assign operator<commit_after>package eval

import (
	"fmt"
	"jacob/black/ast"
	"jacob/black/object"
	"jacob/black/token"
	"math"
)

var (
	// ConstTrue is the only and only true value
	ConstTrue = &object.Boolean{Value: true}
	// ConstFalse is the only and only false value
	ConstFalse = &object.Boolean{Value: false}
	// ConstNil is the only and only nil
	ConstNil = &object.Nil{}
)

func newError(pos token.Position, format string, v ...interface{}) *object.Error {
	return &object.Error{Message: fmt.Sprintf(format, v...), Pos: pos}
}

func isError(o object.Object) bool {
	if o != nil {
		return o.Type() == object.ErrorType
	}
	return false
}

// Eval evaluates the program node and returns an object as a result
func Eval(node ast.Node, env *object.Environment) object.Object {
	switch node := node.(type) {

	// statements
	case *ast.Program:
		return evalProgram(node, env)
	case *ast.LetStatement:
		val := Eval(node.Value, env)
		if isError(val) {
			return val
		}
		env.Set(node.Name.Value, val)
	case *ast.ReturnStatement:
		val := Eval(node.Value, env)
		if isError(val) {
			return val
		}
		return &object.ReturnValue{Value: val}
	case *ast.BlockStatement:
		return evalBlockStatement(node, env)
	case *ast.ExpressionStatement:
		return Eval(node.Expression, env)

		// expressions
	case *ast.PrefixExpression:
		right := Eval(node.Right, env)
		if isError(right) {
			return right
		}
		return evalPrefixExpr(node.Token, right)
	case *ast.InfixExpression:

		if node.Operator == token.Assign {
			if v := evalAssign(node, env); v != nil {
				return v
			}
			return newError(node.Token.Pos, "cannot bind a literal value to a value")
		}

		left := Eval(node.Left, env)
		if isError(left) {
			return left
		}
		right := Eval(node.Right, env)
		if isError(right) {
			return right
		}
		return evalInfixExpr(node.Token, left, right)
	case *ast.IfExpression:
		return evalIfExpr(node, env)

		// literals
	case *ast.IntegerLiteral:
		return &object.Integer{Value: node.Value}
	case *ast.FloatLiteral:
		return &object.Float{Value: node.Value}
	case *ast.BooleanLiteral:
		return boolToBoolean(node.Value)

	case *ast.Identifier:
		return evalIdentifier(node, env)
	}
	return nil
}

func evalProgram(program *ast.Program, env *object.Environment) object.Object {
	var result object.Object

	for _, s := range program.Statements {
		result = Eval(s, env)

		if result != nil {
			// pass up the return type to the top level
			switch result.Type() {
			case object.ReturnType:
				// unwrap the return value
				return result.(*object.ReturnValue).Value
			case object.ErrorType:
				return result.(*object.Error)
			}
		}
	}

	return result
}

func evalBlockStatement(block *ast.BlockStatement, env *object.Environment) object.Object {
	var result object.Object

	for _, s := range block.Statements {
		result = Eval(s, env)

		if result != nil {
			if result.Type() == object.ReturnType || result.Type() == object.ErrorType {
				return result
			}
		}
	}

	return result
}

func evalIfExpr(node *ast.IfExpression, env *object.Environment) object.Object {
	cond := Eval(node.Cond, env)

	if isError(cond) {
		return cond
	}

	if isTruthy(cond) {
		return Eval(node.Do, env)
	} else if node.Else != nil {
		return Eval(node.Else, env)
	}

	return ConstNil
}

// isTruthy - everything is true execpt for false and nil
func isTruthy(o object.Object) bool {
	switch o {
	case ConstFalse, ConstNil:
		return false
	default:
		// special case: 0 or 0.0 is not truthy
		switch o.Type() {
		case object.IntType:
			if o.(*object.Integer).Value == 0 {
				return false
			}
		case object.FloatType:
			if o.(*object.Float).Value == 0.0 {
				return false
			}
		}

		return true
	}
}

func boolToBoolean(b bool) *object.Boolean {
	if b {
		return ConstTrue
	}
	return ConstFalse
}

func evalPrefixExpr(op token.Token, right object.Object) object.Object {
	switch op.Type {
	case token.Bang:
		return evalBangOperatorExpr(right)
	case token.Minus:
		return evalMinusPrefixOperatorExpr(op.Pos, right)
	default:
		return newError(op.Pos, "unknown operator '%s' for type '%s'", op, right.Type())
	}
}

func evalInfixExpr(op token.Token, left object.Object, right object.Object) object.Object {

	if !left.CanApply(op.Type, right.Type()) {
		return newError(op.Pos, "cannot apply operator '%s' for type '%s' and '%s'", op, left.Type(), right.Type())
	}

	// test and convert int to float if needed
	if left.Type() == object.IntType && right.Type() == object.IntType {
		// both ints. easy
		return evalIntegerInfixExpr(op, left, right)
	} else if left.Type() == object.FloatType && right.Type() == object.FloatType {
		// both floats. easy
		return evalFloatInfixExpr(op, left, right)
	}

	// one of them must be a int and the other a float
	if left.Type() == object.IntType || right.Type() == object.IntType {
		if left.Type() == object.FloatType {
			// left is float, right is int.
			// premote right to float
			val := right.(*object.Integer).Value
			right = &object.Float{Value: float64(val)}
			return evalFloatInfixExpr(op, left, right)

		} else if right.Type() == object.FloatType {
			// right is float, left is int
			// premote left to float
			val := left.(*object.Integer).Value
			left = &object.Float{Value: float64(val)}

			return evalFloatInfixExpr(op, left, right)
		}
	}

	// compare actual runtime object
	if op.Type == token.Equal {
		return boolToBoolean(left == right)
	} else if op.Type == token.NotEqual {
		return boolToBoolean(left != right)
	}

	// otherwise 2 objects that don't match
	return newError(op.Pos, "unknown operator '%s' for type '%s' and '%s'", op, left.Type(), right.Type())
}

func evalIntegerInfixExpr(op token.Token, left object.Object, right object.Object) object.Object {
	leftVal := left.(*object.Integer).Value
	rightVal := right.(*object.Integer).Value

	switch op.Type {
	case token.Plus:
		return &object.Integer{Value: leftVal + rightVal}
	case token.Minus:
		return &object.Integer{Value: leftVal - rightVal}
	case token.Times:
		return &object.Integer{Value: leftVal * rightVal}
	case token.Divide:
		if rightVal == 0 {
			return newError(op.Pos, "cannot divide %d by 0", leftVal)
		}
		return &object.Integer{Value: leftVal / rightVal}
	case token.Exp:
		return &object.Integer{Value: int64(math.Pow(float64(leftVal), float64(rightVal)))}
	case token.Mod:
		if rightVal == 0 {
			return newError(op.Pos, "cannot modulo %d by 0", leftVal)
		}
		return &object.Integer{Value: leftVal % rightVal}
	case token.Less:
		return boolToBoolean(leftVal < rightVal)
	case token.Greater:
		return boolToBoolean(leftVal > rightVal)
	case token.Equal:
		return boolToBoolean(leftVal == rightVal)
	case token.NotEqual:
		return boolToBoolean(leftVal != rightVal)
	default:
		return newError(op.Pos, "unknown operator '%s' for type '%s' and '%s'", op.Type, left.Type(), right.Type())
	}
}

func evalFloatInfixExpr(op token.Token, left object.Object, right object.Object) object.Object {
	leftVal := left.(*object.Float).Value
	rightVal := right.(*object.Float).Value

	switch op.Type {
	case token.Plus:
		return &object.Float{Value: leftVal + rightVal}
	case token.Minus:
		return &object.Float{Value: leftVal - rightVal}
	case token.Times:
		return &object.Float{Value: leftVal * rightVal}
	case token.Divide:
		if rightVal == 0 {
			return newError(op.Pos, "cannot divide %d by 0", leftVal)
		}
		return &object.Float{Value: leftVal / rightVal}
	case token.Exp:
		return &object.Float{Value: math.Pow(leftVal, rightVal)}
	case token.Mod:
		if rightVal == 0 {
			return newError(op.Pos, "cannot modulo %d by 0", leftVal)
		}
		return &object.Float{Value: math.Mod(leftVal, rightVal)}
	case token.Less:
		return boolToBoolean(leftVal < rightVal)
	case token.Greater:
		return boolToBoolean(leftVal > rightVal)
	case token.Equal:
		return boolToBoolean(leftVal == rightVal)
	case token.NotEqual:
		return boolToBoolean(leftVal != rightVal)
	default:
		return newError(op.Pos, "unknown operator '%s' for type '%s' and '%s'", op.Type, left.Type(), right.Type())
	}
}

func evalBangOperatorExpr(right object.Object) object.Object {
	switch right {
	case ConstTrue:
		return ConstFalse
	case ConstFalse:
		return ConstTrue
	case ConstNil:
		return ConstTrue
	default:
		return boolToBoolean(!isTruthy(right))
	}
}

func evalMinusPrefixOperatorExpr(pos token.Position, right object.Object) object.Object {
	switch right.Type() {
	case object.IntType:
		v := right.(*object.Integer).Value
		return &object.Integer{Value: -v}
	case object.FloatType:
		v := right.(*object.Float).Value
		return &object.Float{Value: -v}
	default:
		return newError(pos, "unknown operator '-' for type '%s'", right.Type())
	}
}

func evalIdentifier(id *ast.Identifier, env *object.Environment) object.Object {
	if val, ok := env.Get(id.Value); ok {
		return val
	}
	return newError(id.Token.Pos, "identifier not found: %s", id.Value)
}

func evalAssign(node *ast.InfixExpression, env *object.Environment) object.Object {
	// special case = assign operator
	switch l := node.Left.(type) {
	case *ast.Identifier:
		// check if exists already
		if val, ok := env.Get(l.Value); ok {
			if isError(val) {
				return val
			}

			// eval rhs
			right := Eval(node.Right, env)
			if isError(right) {
				return right
			}

			// must be same type
			if val.Type() == right.Type() {
				env.Set(l.Value, right)
				return right
			}

			return newError(l.Token.Pos, "cannot assign variable '%s' of type '%s' to value '%s' of type '%s'", l.Value, val.Type(), right, right.Type())
		}
		return newError(l.Token.Pos, "cannot assign value to variable '%s' that does not exist", l.Value)
	default:
		return nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	proto         = "tcp"
	retryInterval = 1
)

var (
	listeningAddress    = flag.String("listeningAddress", ":8080", "Address on which to expose Prometheus metrics.")
	muninAddress        = flag.String("muninAddress", "localhost:4949", "munin-node address.")
	muninScrapeInterval = flag.Int("muninScrapeInterval", 60, "Interval in seconds between scrapes.")
	globalConn          net.Conn
	hostname            string
	graphs              []string
	gaugePerMetric      map[string]*prometheus.GaugeVec
	counterPerMetric     map[string]*prometheus.CounterVec
	muninBanner         *regexp.Regexp
)

func init() {
	flag.Parse()
	var err error
	gaugePerMetric = map[string]*prometheus.GaugeVec{}
	counterPerMetric = map[string]*prometheus.CounterVec{}
	muninBanner = regexp.MustCompile(`# munin node at (.*)`)

	err = connect()
	if err != nil {
		log.Fatalf("Could not connect to %s: %s", *muninAddress, err)
	}
}

func serveStatus() {
	http.Handle("/metrics", prometheus.Handler())
	http.ListenAndServe(*listeningAddress, nil)
}

func connect() (err error) {
	log.Printf("Connecting...")
	globalConn, err = net.Dial(proto, *muninAddress)
	if err != nil {
		return
	}
	log.Printf("connected!")

	reader := bufio.NewReader(globalConn)
	head, err := reader.ReadString('\n')
	if err != nil {
		return
	}

	matches := muninBanner.FindStringSubmatch(head)
	if len(matches) != 2 { // expect: # munin node at <hostname>
		return fmt.Errorf("Unexpected line: %s", head)
	}
	hostname = matches[1]
	log.Printf("Found hostname: %s", hostname)
	return
}

func muninCommand(cmd string) (reader *bufio.Reader, err error) {
	reader = bufio.NewReader(globalConn)

	fmt.Fprintf(globalConn, cmd+"\n")

	_, err = reader.Peek(1)
	switch err {
	case io.EOF:
		log.Printf("not connected anymore, closing connection")
		globalConn.Close()
		for {
			err = connect()
			if err == nil {
				break
			}
			log.Printf("Couldn't reconnect: %s", err)
			time.Sleep(retryInterval * time.Second)
		}

		return muninCommand(cmd)
	case nil: //no error
		break
	default:
		log.Fatalf("Unexpected error: %s", err)
	}

	return
}

func muninList() (items []string, err error) {
	munin, err := muninCommand("list")
	if err != nil {
		log.Printf("couldn't get list")
		return
	}

	response, err := munin.ReadString('\n') // we are only interested in the first line
	if err != nil {
		log.Printf("couldn't read response")
		return
	}

	if response[0] == '#' { // # not expected here
		err = fmt.Errorf("Error getting items: %s", response)
		return
	}
	items = strings.Fields(strings.TrimRight(response, "\n"))
	return
}

func muninConfig(name string) (config map[string]map[string]string, graphConfig map[string]string, err error) {
	graphConfig = make(map[string]string)
	config = make(map[string]map[string]string)

	resp, err := muninCommand("config " + name)
	if err != nil {
		log.Printf("couldn't get config for %s", name)
		return
	}

	for {
		line, err := resp.ReadString('\n')
		if err == io.EOF {
			log.Fatalf("unexpected EOF, retrying")
			return muninConfig(name)
		}
		if err != nil {
			return nil, nil, err
		}
		if line == ".\n" { // munin end marker
			break
		}
		if line[0] == '#' { // here it's just a comment, so ignore it
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 2 {
			return nil, nil, fmt.Errorf("Line unexpected: %s", line)
		}
		key, value := parts[0], strings.TrimRight(strings.Join(parts[1:], " "), "\n")

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 { // it's a metric config (metric.label etc)
			if _, ok := config[keyParts[0]]; !ok { //FIXME: is there no better way?
				config[keyParts[0]] = make(map[string]string)
			}
			config[keyParts[0]][keyParts[1]] = value
		} else {
			graphConfig[keyParts[0]] = value
		}
	}
	return
}

func registerMetrics() (err error) {
	items, err := muninList()
	if err != nil {
		return
	}

	for _, name := range items {
		graphs = append(graphs, name)
		configs, graphConfig, err := muninConfig(name)
		if err != nil {
			return err
		}

		for metric, config := range configs {
			metricName := strings.Replace(name + "_" + metric, "-","_",-1)
			desc := graphConfig["graph_title"] + ": " + config["label"]
			if config["info"] != "" {
				desc = desc + ", " + config["info"]
			}
			muninType := strings.ToLower(config["type"])
			// muninType can be empty and defaults to gauge
			if muninType == "counter" || muninType == "derive" {
	                        gv := prometheus.NewCounterVec(
        	                        prometheus.CounterOpts{
                	                        Name: metricName,
                        	                Help: desc,
						ConstLabels: prometheus.Labels{"type":muninType},
                                	},
                                	[]string{"hostname","graphname","muninlabel"},
                        	)
				log.Printf("Registered counter %s: %s", metricName, desc)
                        	counterPerMetric[metricName] = gv
                        	prometheus.Register(gv)

			} else {
                        	gv := prometheus.NewGaugeVec(
                                	prometheus.GaugeOpts{
                                        	Name: metricName,
	                                        Help: desc,
						ConstLabels: prometheus.Labels{"type":"counter"},
                	                },
                        	        []string{"hostname","graphname","muninlabel"},
                        	)
				log.Printf("Registered gauge %s: %s", metricName, desc)
        	                gaugePerMetric[metricName] = gv
                	        prometheus.Register(gv)
			}
		}
	}
	return nil
}

func fetchMetrics() (err error) {
	for _, graph := range graphs {
		munin, err := muninCommand("fetch " + graph)
		if err != nil {
			return err
		}

		for {
			line, err := munin.ReadString('\n')
			line = strings.TrimRight(line, "\n")
			if err == io.EOF {
				log.Fatalf("unexpected EOF, retrying")
				return fetchMetrics()
			}
			if err != nil {
				return err
			}
			if len(line) == 1 && line[0] == '.' {
				log.Printf("End of list")
				break
			}

			parts := strings.Fields(line)
			if len(parts) != 2 {
				log.Printf("unexpected line: %s", line)
				continue
			}
			key, valueString := strings.Split(parts[0], ".")[0], parts[1]
			value, err := strconv.ParseFloat(valueString, 64)
			if err != nil {
				log.Printf("Couldn't parse value in line %s, malformed?", line)
				continue
			}
			name := strings.Replace(graph + "_" + key, "-","_",-1)
			log.Printf("%s: %f\n", name, value)
			_, isGauge := gaugePerMetric[name]
			if isGauge {
	                        gaugePerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			} else {
				counterPerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := registerMetrics()
	if err != nil {
		log.Fatalf("Could not register metrics: %s", err)
	}

	go serveStatus()

	func() {
		for {
			log.Printf("Scraping")
			err := fetchMetrics()
			if err != nil {
				log.Printf("Error occured when trying to fetch metrics: %s", err)
			}
			time.Sleep(time.Duration(*muninScrapeInterval) * time.Second)
		}
	}()
}
<commit_msg>Stupid error, confusing counter/gauge, at least just in the label.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	proto         = "tcp"
	retryInterval = 1
)

var (
	listeningAddress    = flag.String("listeningAddress", ":8080", "Address on which to expose Prometheus metrics.")
	muninAddress        = flag.String("muninAddress", "localhost:4949", "munin-node address.")
	muninScrapeInterval = flag.Int("muninScrapeInterval", 60, "Interval in seconds between scrapes.")
	globalConn          net.Conn
	hostname            string
	graphs              []string
	gaugePerMetric      map[string]*prometheus.GaugeVec
	counterPerMetric     map[string]*prometheus.CounterVec
	muninBanner         *regexp.Regexp
)

func init() {
	flag.Parse()
	var err error
	gaugePerMetric = map[string]*prometheus.GaugeVec{}
	counterPerMetric = map[string]*prometheus.CounterVec{}
	muninBanner = regexp.MustCompile(`# munin node at (.*)`)

	err = connect()
	if err != nil {
		log.Fatalf("Could not connect to %s: %s", *muninAddress, err)
	}
}

func serveStatus() {
	http.Handle("/metrics", prometheus.Handler())
	http.ListenAndServe(*listeningAddress, nil)
}

func connect() (err error) {
	log.Printf("Connecting...")
	globalConn, err = net.Dial(proto, *muninAddress)
	if err != nil {
		return
	}
	log.Printf("connected!")

	reader := bufio.NewReader(globalConn)
	head, err := reader.ReadString('\n')
	if err != nil {
		return
	}

	matches := muninBanner.FindStringSubmatch(head)
	if len(matches) != 2 { // expect: # munin node at <hostname>
		return fmt.Errorf("Unexpected line: %s", head)
	}
	hostname = matches[1]
	log.Printf("Found hostname: %s", hostname)
	return
}

func muninCommand(cmd string) (reader *bufio.Reader, err error) {
	reader = bufio.NewReader(globalConn)

	fmt.Fprintf(globalConn, cmd+"\n")

	_, err = reader.Peek(1)
	switch err {
	case io.EOF:
		log.Printf("not connected anymore, closing connection")
		globalConn.Close()
		for {
			err = connect()
			if err == nil {
				break
			}
			log.Printf("Couldn't reconnect: %s", err)
			time.Sleep(retryInterval * time.Second)
		}

		return muninCommand(cmd)
	case nil: //no error
		break
	default:
		log.Fatalf("Unexpected error: %s", err)
	}

	return
}

func muninList() (items []string, err error) {
	munin, err := muninCommand("list")
	if err != nil {
		log.Printf("couldn't get list")
		return
	}

	response, err := munin.ReadString('\n') // we are only interested in the first line
	if err != nil {
		log.Printf("couldn't read response")
		return
	}

	if response[0] == '#' { // # not expected here
		err = fmt.Errorf("Error getting items: %s", response)
		return
	}
	items = strings.Fields(strings.TrimRight(response, "\n"))
	return
}

func muninConfig(name string) (config map[string]map[string]string, graphConfig map[string]string, err error) {
	graphConfig = make(map[string]string)
	config = make(map[string]map[string]string)

	resp, err := muninCommand("config " + name)
	if err != nil {
		log.Printf("couldn't get config for %s", name)
		return
	}

	for {
		line, err := resp.ReadString('\n')
		if err == io.EOF {
			log.Fatalf("unexpected EOF, retrying")
			return muninConfig(name)
		}
		if err != nil {
			return nil, nil, err
		}
		if line == ".\n" { // munin end marker
			break
		}
		if line[0] == '#' { // here it's just a comment, so ignore it
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 2 {
			return nil, nil, fmt.Errorf("Line unexpected: %s", line)
		}
		key, value := parts[0], strings.TrimRight(strings.Join(parts[1:], " "), "\n")

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 { // it's a metric config (metric.label etc)
			if _, ok := config[keyParts[0]]; !ok { //FIXME: is there no better way?
				config[keyParts[0]] = make(map[string]string)
			}
			config[keyParts[0]][keyParts[1]] = value
		} else {
			graphConfig[keyParts[0]] = value
		}
	}
	return
}

func registerMetrics() (err error) {
	items, err := muninList()
	if err != nil {
		return
	}

	for _, name := range items {
		graphs = append(graphs, name)
		configs, graphConfig, err := muninConfig(name)
		if err != nil {
			return err
		}

		for metric, config := range configs {
			metricName := strings.Replace(name + "_" + metric, "-","_",-1)
			desc := graphConfig["graph_title"] + ": " + config["label"]
			if config["info"] != "" {
				desc = desc + ", " + config["info"]
			}
			muninType := strings.ToLower(config["type"])
			// muninType can be empty and defaults to gauge
			if muninType == "counter" || muninType == "derive" {
	                        gv := prometheus.NewCounterVec(
        	                        prometheus.CounterOpts{
                	                        Name: metricName,
                        	                Help: desc,
						ConstLabels: prometheus.Labels{"type":muninType},
                                	},
                                	[]string{"hostname","graphname","muninlabel"},
                        	)
				log.Printf("Registered counter %s: %s", metricName, desc)
                        	counterPerMetric[metricName] = gv
                        	prometheus.Register(gv)

			} else {
                        	gv := prometheus.NewGaugeVec(
                                	prometheus.GaugeOpts{
                                        	Name: metricName,
	                                        Help: desc,
						ConstLabels: prometheus.Labels{"type":"gauge"},
                	                },
                        	        []string{"hostname","graphname","muninlabel"},
                        	)
				log.Printf("Registered gauge %s: %s", metricName, desc)
        	                gaugePerMetric[metricName] = gv
                	        prometheus.Register(gv)
			}
		}
	}
	return nil
}

func fetchMetrics() (err error) {
	for _, graph := range graphs {
		munin, err := muninCommand("fetch " + graph)
		if err != nil {
			return err
		}

		for {
			line, err := munin.ReadString('\n')
			line = strings.TrimRight(line, "\n")
			if err == io.EOF {
				log.Fatalf("unexpected EOF, retrying")
				return fetchMetrics()
			}
			if err != nil {
				return err
			}
			if len(line) == 1 && line[0] == '.' {
				log.Printf("End of list")
				break
			}

			parts := strings.Fields(line)
			if len(parts) != 2 {
				log.Printf("unexpected line: %s", line)
				continue
			}
			key, valueString := strings.Split(parts[0], ".")[0], parts[1]
			value, err := strconv.ParseFloat(valueString, 64)
			if err != nil {
				log.Printf("Couldn't parse value in line %s, malformed?", line)
				continue
			}
			name := strings.Replace(graph + "_" + key, "-","_",-1)
			log.Printf("%s: %f\n", name, value)
			_, isGauge := gaugePerMetric[name]
			if isGauge {
	                        gaugePerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			} else {
				counterPerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := registerMetrics()
	if err != nil {
		log.Fatalf("Could not register metrics: %s", err)
	}

	go serveStatus()

	func() {
		for {
			log.Printf("Scraping")
			err := fetchMetrics()
			if err != nil {
				log.Printf("Error occured when trying to fetch metrics: %s", err)
			}
			time.Sleep(time.Duration(*muninScrapeInterval) * time.Second)
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar_test

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	clientsBarBar "github.com/uber/zanzibar/examples/example-gateway/build/gen-code/clients/bar/bar"
	exampleGateway "github.com/uber/zanzibar/examples/example-gateway/build/services/example-gateway"
	"github.com/uber/zanzibar/runtime"
	"github.com/uber/zanzibar/test/lib/bench_gateway"
	"github.com/uber/zanzibar/test/lib/test_gateway"
	"github.com/uber/zanzibar/test/lib/util"
)

var defaultTestOptions *testGateway.Options = &testGateway.Options{
	KnownHTTPBackends:     []string{"bar", "contacts", "google-now"},
	KnownTChannelBackends: []string{"baz"},
	ConfigFiles:           util.DefaultConfigFiles("example-gateway"),
}
var defaultTestConfig map[string]interface{} = map[string]interface{}{
	"clients.baz.serviceName": "baz",
}

func TestMakingClientWriteJSONWithBadJSON(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)
	client := zanzibar.NewHTTPClient(
		bgateway.ActualGateway.Logger,
		bgateway.ActualGateway.AllHostScope,
		"clientID",
		[]string{"DoStuff"},
		"/",
		map[string]string{},
		time.Second,
	)
	req := zanzibar.NewClientHTTPRequest("clientID", "DoStuff", client)

	err = req.WriteJSON("GET", "/foo", nil, &failingJsonObj{})
	assert.NotNil(t, err)
	assert.Equal(t,
		"Could not serialize clientID.DoStuff request json: cannot serialize",
		err.Error(),
	)

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Could not serialize request json"], 1)
}

func TestMakingClientWriteJSONWithBadHTTPMethod(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)
	client := zanzibar.NewHTTPClient(
		bgateway.ActualGateway.Logger,
		bgateway.ActualGateway.AllHostScope,
		"clientID",
		[]string{"DoStuff"},
		"/",
		map[string]string{},
		time.Second,
	)
	req := zanzibar.NewClientHTTPRequest("clientID", "DoStuff", client)

	err = req.WriteJSON("@INVALIDMETHOD", "/foo", nil, nil)
	assert.NotNil(t, err)
	assert.Equal(t,
		"Could not create outbound clientID.DoStuff request: net/http: invalid method \"@INVALIDMETHOD\"",
		err.Error(),
	)

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Could not create outbound request"], 1)
}

func TestMakingClientCalLWithHeaders(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(200)
			_, _ = w.Write([]byte(r.Header.Get("Example-Header")))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	barClient := deps.Client.Bar
	client := barClient.HTTPClient()

	req := zanzibar.NewClientHTTPRequest("bar", "Normal", client)

	err = req.WriteJSON(
		"POST",
		client.BaseURL+"/bar-path",
		map[string]string{
			"Example-Header": "Example-Value",
		},
		nil,
	)
	assert.NoError(t, err)

	res, err := req.Do(context.Background())
	assert.NoError(t, err)
	assert.Equal(t, 200, res.StatusCode)

	bytes, err := res.ReadAll()
	assert.NoError(t, err)
	assert.Equal(t, []byte("Example-Value"), bytes)

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestBarClientWithoutHeaders(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bar := deps.Client.Bar

	_, _, err = bar.EchoI8(
		context.Background(), nil, &clientsBarBar.Echo_EchoI8_Args{Arg: 42},
	)

	assert.NotNil(t, err)
	assert.Equal(t, "Missing mandatory header: x-uuid", err.Error())

	logs := gateway.AllLogs()

	assert.Equal(t, 1, len(logs))

	lines := logs["Got outbound request without mandatory header"]
	assert.Equal(t, 1, len(lines))

	logLine := lines[0]
	assert.Equal(t, "bar", logLine["clientID"])
	assert.Equal(t, "EchoI8", logLine["methodName"])
	assert.Equal(t, "x-uuid", logLine["headerName"])
}

func TestMakingClientCallWithRespHeaders(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Add("Example-Header", "Example-Value")
			w.WriteHeader(200)
			_, _ = w.Write([]byte(`{
				"stringField":"foo",
				"intWithRange": 0,
				"intWithoutRange": 1,
				"mapIntWithRange": {},
				"mapIntWithoutRange": {},
				"binaryField": "d29ybGQ="
			}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, headers, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.NoError(t, err)
	assert.NotNil(t, body)
	assert.Equal(t, "Example-Value", headers["Example-Header"])

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingClientCallWithThriftException(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(403)
			_, _ = w.Write([]byte(`{"stringField":"test"}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, _, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.Error(t, err)
	assert.Nil(t, body)

	realError := err.(*clientsBarBar.BarException)
	assert.Equal(t, realError.StringField, "test")

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingClientCallWithBadStatusCode(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(402)
			_, _ = w.Write([]byte(`{"stringField":"test"}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, _, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.Error(t, err)
	assert.Nil(t, body)
	assert.Equal(t, "Unexpected http client response (402)", err.Error())

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Unknown response status code"], 1)
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingCallWithThriftException(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	gateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/arg-not-struct-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(403)
			_, _ = w.Write([]byte(`{"stringField":"test"}`))
		},
	)

	bgateway := gateway.(*benchGateway.BenchGateway)
	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	_, err = bClient.ArgNotStruct(
		context.Background(), nil,
		&clientsBarBar.Bar_ArgNotStruct_Args{
			Request: "request",
		},
	)
	assert.Error(t, err)

	realError := err.(*clientsBarBar.BarException)
	assert.Equal(t, realError.StringField, "test")

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingClientCallWithServerError(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(500)
			_, _ = w.Write([]byte(`{}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, _, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.Error(t, err)
	assert.Nil(t, body)
	assert.Equal(t, "Unexpected http client response (500)", err.Error())

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Unknown response status code"], 1)
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}
<commit_msg>Test that default headers actually get sent<commit_after>// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar_test

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	clientsBarBar "github.com/uber/zanzibar/examples/example-gateway/build/gen-code/clients/bar/bar"
	exampleGateway "github.com/uber/zanzibar/examples/example-gateway/build/services/example-gateway"
	"github.com/uber/zanzibar/runtime"
	"github.com/uber/zanzibar/test/lib/bench_gateway"
	"github.com/uber/zanzibar/test/lib/test_gateway"
	"github.com/uber/zanzibar/test/lib/util"
)

var defaultTestOptions *testGateway.Options = &testGateway.Options{
	KnownHTTPBackends:     []string{"bar", "contacts", "google-now"},
	KnownTChannelBackends: []string{"baz"},
	ConfigFiles:           util.DefaultConfigFiles("example-gateway"),
}
var defaultTestConfig map[string]interface{} = map[string]interface{}{
	"clients.baz.serviceName": "baz",
}

func TestMakingClientWriteJSONWithBadJSON(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)
	client := zanzibar.NewHTTPClient(
		bgateway.ActualGateway.Logger,
		bgateway.ActualGateway.AllHostScope,
		"clientID",
		[]string{"DoStuff"},
		"/",
		map[string]string{},
		time.Second,
	)
	req := zanzibar.NewClientHTTPRequest("clientID", "DoStuff", client)

	err = req.WriteJSON("GET", "/foo", nil, &failingJsonObj{})
	assert.NotNil(t, err)
	assert.Equal(t,
		"Could not serialize clientID.DoStuff request json: cannot serialize",
		err.Error(),
	)

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Could not serialize request json"], 1)
}

func TestMakingClientWriteJSONWithBadHTTPMethod(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)
	client := zanzibar.NewHTTPClient(
		bgateway.ActualGateway.Logger,
		bgateway.ActualGateway.AllHostScope,
		"clientID",
		[]string{"DoStuff"},
		"/",
		map[string]string{},
		time.Second,
	)
	req := zanzibar.NewClientHTTPRequest("clientID", "DoStuff", client)

	err = req.WriteJSON("@INVALIDMETHOD", "/foo", nil, nil)
	assert.NotNil(t, err)
	assert.Equal(t,
		"Could not create outbound clientID.DoStuff request: net/http: invalid method \"@INVALIDMETHOD\"",
		err.Error(),
	)

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Could not create outbound request"], 1)
}

func TestMakingClientCallWithHeaders(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(200)
			_, _ = w.Write([]byte(r.Header.Get("Example-Header")))
			// Check that the default header got set and actually sent to the server.
			assert.Equal(t, r.Header.Get("X-Client-ID"), "bar")
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	barClient := deps.Client.Bar
	client := barClient.HTTPClient()

	req := zanzibar.NewClientHTTPRequest("bar", "Normal", client)

	err = req.WriteJSON(
		"POST",
		client.BaseURL+"/bar-path",
		map[string]string{
			"Example-Header": "Example-Value",
		},
		nil,
	)
	assert.NoError(t, err)

	res, err := req.Do(context.Background())
	assert.NoError(t, err)
	assert.Equal(t, 200, res.StatusCode)

	bytes, err := res.ReadAll()
	assert.NoError(t, err)
	assert.Equal(t, []byte("Example-Value"), bytes)

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestBarClientWithoutHeaders(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bar := deps.Client.Bar

	_, _, err = bar.EchoI8(
		context.Background(), nil, &clientsBarBar.Echo_EchoI8_Args{Arg: 42},
	)

	assert.NotNil(t, err)
	assert.Equal(t, "Missing mandatory header: x-uuid", err.Error())

	logs := gateway.AllLogs()

	assert.Equal(t, 1, len(logs))

	lines := logs["Got outbound request without mandatory header"]
	assert.Equal(t, 1, len(lines))

	logLine := lines[0]
	assert.Equal(t, "bar", logLine["clientID"])
	assert.Equal(t, "EchoI8", logLine["methodName"])
	assert.Equal(t, "x-uuid", logLine["headerName"])
}

func TestMakingClientCallWithRespHeaders(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Add("Example-Header", "Example-Value")
			w.WriteHeader(200)
			_, _ = w.Write([]byte(`{
				"stringField":"foo",
				"intWithRange": 0,
				"intWithoutRange": 1,
				"mapIntWithRange": {},
				"mapIntWithoutRange": {},
				"binaryField": "d29ybGQ="
			}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, headers, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.NoError(t, err)
	assert.NotNil(t, body)
	assert.Equal(t, "Example-Value", headers["Example-Header"])

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingClientCallWithThriftException(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(403)
			_, _ = w.Write([]byte(`{"stringField":"test"}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, _, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.Error(t, err)
	assert.Nil(t, body)

	realError := err.(*clientsBarBar.BarException)
	assert.Equal(t, realError.StringField, "test")

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingClientCallWithBadStatusCode(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(402)
			_, _ = w.Write([]byte(`{"stringField":"test"}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, _, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.Error(t, err)
	assert.Nil(t, body)
	assert.Equal(t, "Unexpected http client response (402)", err.Error())

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Unknown response status code"], 1)
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingCallWithThriftException(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	gateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/arg-not-struct-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(403)
			_, _ = w.Write([]byte(`{"stringField":"test"}`))
		},
	)

	bgateway := gateway.(*benchGateway.BenchGateway)
	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	_, err = bClient.ArgNotStruct(
		context.Background(), nil,
		&clientsBarBar.Bar_ArgNotStruct_Args{
			Request: "request",
		},
	)
	assert.Error(t, err)

	realError := err.(*clientsBarBar.BarException)
	assert.Equal(t, realError.StringField, "test")

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}

func TestMakingClientCallWithServerError(t *testing.T) {
	gateway, err := benchGateway.CreateGateway(
		defaultTestConfig,
		defaultTestOptions,
		exampleGateway.CreateGateway,
	)
	if !assert.NoError(t, err) {
		return
	}
	defer gateway.Close()

	bgateway := gateway.(*benchGateway.BenchGateway)

	bgateway.HTTPBackends()["bar"].HandleFunc(
		"POST", "/bar-path",
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(500)
			_, _ = w.Write([]byte(`{}`))
		},
	)

	deps := bgateway.Dependencies.(*exampleGateway.DependenciesTree)
	bClient := deps.Client.Bar

	body, _, err := bClient.Normal(
		context.Background(), nil, &clientsBarBar.Bar_Normal_Args{},
	)
	assert.Error(t, err)
	assert.Nil(t, body)
	assert.Equal(t, "Unexpected http client response (500)", err.Error())

	logs := bgateway.AllLogs()
	assert.Len(t, logs["Unknown response status code"], 1)
	assert.Len(t, logs["Finished an outgoing client HTTP request"], 1)
}
<|endoftext|>


<commit_before>package go_simple_sql

import (
	"database/sql"
	"fmt"
)

type CONN struct {
	DB sql.DB
}

func (c *CONN) InitDB(ip, port, user, pwd, dbname, charset string) error {
	url := user + ":" + pwd + "@" + "tcp(" + ip + ":" + port + ")/" + dbname + "?charset=" + charset
	db, err := sql.Open("mysql", url)
	if err != nil {
		fmt.Println("mysql init fail")
	} else {
		c.DB = *db
		fmt.Println("mysql init success")
	}
}

func (c *CONN) Query(text string) ([]map[string]string, error) {
	rows, err := c.DB.Query(text)
	result := make([]map[string]string, 0)
	if err != nil {
		return result, err
	}
	columns, _ := rows.Columns()
	scanArgs := make([]interface{}, len(columns))
	values := make([]interface{}, len(columns))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	for rows.Next() {
		err = rows.Scan(scanArgs...)
		record := make(map[string]string)
		for i, col := range values {
			if col != nil {
				record[columns[i]] = string(col.([]byte))
			}
		}
		result = append(result, record)
	}
	return result, err
}

func (c *CONN) Update(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}

func (c *CONN) Insert(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	id, err := result.LastInsertId()
	return id, err
}

func (c *CONN) Delete(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}
<commit_msg>fix bug<commit_after>package go_simple_sql

import (
	"database/sql"
	"fmt"
)

type CONN struct {
	DB sql.DB
}

func (c *CONN) InitDB(ip, port, user, pwd, dbname, charset string) error {
	url := user + ":" + pwd + "@" + "tcp(" + ip + ":" + port + ")/" + dbname + "?charset=" + charset
	db, err := sql.Open("mysql", url)
	if err != nil {
		fmt.Println("mysql init fail")
	} else {
		c.DB = *db
		fmt.Println("mysql init success")
	}
	return err
}

func (c *CONN) Query(text string) ([]map[string]string, error) {
	rows, err := c.DB.Query(text)
	result := make([]map[string]string, 0)
	if err != nil {
		return result, err
	}
	columns, _ := rows.Columns()
	scanArgs := make([]interface{}, len(columns))
	values := make([]interface{}, len(columns))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	for rows.Next() {
		err = rows.Scan(scanArgs...)
		record := make(map[string]string)
		for i, col := range values {
			if col != nil {
				record[columns[i]] = string(col.([]byte))
			}
		}
		result = append(result, record)
	}
	return result, err
}

func (c *CONN) Update(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}

func (c *CONN) Insert(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	id, err := result.LastInsertId()
	return id, err
}

func (c *CONN) Delete(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}
<|endoftext|>


<commit_before>/****************************************************************************
 * This file is part of Builder.
 *
 * Copyright (C) 2015 Pier Luigi Fiorini
 *
 * Author(s):
 *    Pier Luigi Fiorini <pierluigi.fiorini@gmail.com>
 *
 * $BEGIN_LICENSE:AGPL3+$
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * $END_LICENSE$
 ***************************************************************************/

package database

import (
	"bytes"
	"encoding/json"
	"github.com/boltdb/bolt"
)

type Image struct {
	Name          string   `json:"name"`
	Description   string   `json:"descr"`
	Architectures []string `json:"archs"`
	Vcs           VcsInfo  `json:"vcs"`
}

// Return whether the image was stored into the db.
func (db *Database) HasImage(name string) bool {
	var found bool = false
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, _ := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, _ = c.Next() {
			found = true
			return nil
		}
		return nil
	})
	return found
}

// Return a list of image names.
func (db *Database) GetImageNames() []string {
	var list = []string{}
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			list = append(list, string(k))
			return nil
		})
		return nil
	})
	return list
}

// Return a list of all images.
func (db *Database) ListAllImages() []*Image {
	var list []*Image
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			img := &Image{}
			json.Unmarshal(v, &img)
			list = append(list, img)
			return nil
		})
		return nil
	})
	return list
}

// Return an image from the database.
func (db *Database) GetImage(name string) *Image {
	var img *Image = nil
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, v := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, v = c.Next() {
			json.Unmarshal(v, &img)
			return nil
		}
		return nil
	})
	return img
}

// Add an image to the database.
func (db *Database) AddImage(img *Image) error {
	encoded, err := json.Marshal(img)
	if err != nil {
		return nil
	}

	err = db.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte("image"))
		if err != nil {
			return err
		}

		err = bucket.Put([]byte(img.Name), encoded)
		if err != nil {
			return err
		}

		return nil
	})

	if err == nil {
		// Add the architectures as supported
		err = db.SaveArchitectures(img.Architectures...)
		if err != nil {
			return err
		}
	}

	return err
}

// Remove an image from the database.
func (db *Database) RemoveImage(name string) error {
	return db.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		// Unmarshal the image
		img := &Image{}
		c := bucket.Cursor()
		for k, v := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, v = c.Next() {
			if err := json.Unmarshal(v, &img); err != nil {
				return err
			}
		}

		// Delete the bucket
		err := bucket.Delete([]byte(name))
		if err != nil {
			return err
		}

		// Remove these architectures if they are not referenced
		// by any other package or image
		if err := db.RemoveArchitectures(img.Architectures...); err != nil {
			return err
		}

		return nil
	})
}

// Iterate the images list.
func (db *Database) ForEachImage(f func(img *Image)) {
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			var img *Image
			err := json.Unmarshal(v, &img)
			if err != nil {
				return err
			}

			f(img)

			return nil
		})

		return nil
	})
}
<commit_msg>Return an error if the image cannot be serialized<commit_after>/****************************************************************************
 * This file is part of Builder.
 *
 * Copyright (C) 2015 Pier Luigi Fiorini
 *
 * Author(s):
 *    Pier Luigi Fiorini <pierluigi.fiorini@gmail.com>
 *
 * $BEGIN_LICENSE:AGPL3+$
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * $END_LICENSE$
 ***************************************************************************/

package database

import (
	"bytes"
	"encoding/json"
	"github.com/boltdb/bolt"
)

type Image struct {
	Name          string   `json:"name"`
	Description   string   `json:"descr"`
	Architectures []string `json:"archs"`
	Vcs           VcsInfo  `json:"vcs"`
}

// Return whether the image was stored into the db.
func (db *Database) HasImage(name string) bool {
	var found bool = false
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, _ := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, _ = c.Next() {
			found = true
			return nil
		}
		return nil
	})
	return found
}

// Return a list of image names.
func (db *Database) GetImageNames() []string {
	var list = []string{}
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			list = append(list, string(k))
			return nil
		})
		return nil
	})
	return list
}

// Return a list of all images.
func (db *Database) ListAllImages() []*Image {
	var list []*Image
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			img := &Image{}
			json.Unmarshal(v, &img)
			list = append(list, img)
			return nil
		})
		return nil
	})
	return list
}

// Return an image from the database.
func (db *Database) GetImage(name string) *Image {
	var img *Image = nil
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, v := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, v = c.Next() {
			json.Unmarshal(v, &img)
			return nil
		}
		return nil
	})
	return img
}

// Add an image to the database.
func (db *Database) AddImage(img *Image) error {
	encoded, err := json.Marshal(img)
	if err != nil {
		return err
	}

	err = db.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte("image"))
		if err != nil {
			return err
		}

		err = bucket.Put([]byte(img.Name), encoded)
		if err != nil {
			return err
		}

		return nil
	})

	if err == nil {
		// Add the architectures as supported
		err = db.SaveArchitectures(img.Architectures...)
		if err != nil {
			return err
		}
	}

	return err
}

// Remove an image from the database.
func (db *Database) RemoveImage(name string) error {
	return db.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		// Unmarshal the image
		img := &Image{}
		c := bucket.Cursor()
		for k, v := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, v = c.Next() {
			if err := json.Unmarshal(v, &img); err != nil {
				return err
			}
		}

		// Delete the bucket
		err := bucket.Delete([]byte(name))
		if err != nil {
			return err
		}

		// Remove these architectures if they are not referenced
		// by any other package or image
		if err := db.RemoveArchitectures(img.Architectures...); err != nil {
			return err
		}

		return nil
	})
}

// Iterate the images list.
func (db *Database) ForEachImage(f func(img *Image)) {
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("image"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			var img *Image
			err := json.Unmarshal(v, &img)
			if err != nil {
				return err
			}

			f(img)

			return nil
		})

		return nil
	})
}
<|endoftext|>


<commit_before>package processors

import (
	"fmt"
	"time"

	"github.com/jcelliott/lumber"
	"github.com/nanobox-io/golang-docker-client"

	process_provider "github.com/nanobox-io/nanobox/processors/provider"
	"github.com/nanobox-io/nanobox/util"
	"github.com/nanobox-io/nanobox/util/display"
	"github.com/nanobox-io/nanobox/util/provider"
	//	"github.com/nanobox-io/nanobox/util/update"
)

func Update() error {

	// init docker client
	if err := process_provider.Init(); err != nil {
		return fmt.Errorf("failed to init docker client: %s", err.Error())
	}

	// // check to see if nanobox needs to update
	// update.Check()

	// update all the nanobox images
	pullImages()

	// pull the latest docker-machine image
	return provider.Install()
}

func pullImages() error {
	display.OpenContext("Updating Images")
	defer display.CloseContext()

	images, err := docker.ImageList()
	if err != nil {
		return err
	}

	for _, image := range images {
		display.StartTask("Pulling %s image", image.Slug)

		// generate a docker percent display
		dockerPercent := &display.DockerPercentDisplay{
			Output: display.NewStreamer("info"),
			// Prefix: buildImage,
		}

		// pull the build image
		imagePullFunc := func() error {
			_, err := docker.ImagePull(image.Slug, dockerPercent)
			return err
		}

		if err := util.Retry(imagePullFunc, 5, time.Second); err != nil {
			lumber.Error("code:pullBuildImage:docker.ImagePull(%s, nil): %s", image.Slug, err.Error())
			display.ErrorTask()
			return fmt.Errorf("failed to pull docker image (%s): %s", image.Slug, err.Error())
		}

		display.StopTask()
	}

	return nil
}
<commit_msg>hotfix/imagepull<commit_after>package processors

import (
	"fmt"
	"time"

	"github.com/jcelliott/lumber"
	"github.com/nanobox-io/golang-docker-client"

	process_provider "github.com/nanobox-io/nanobox/processors/provider"
	"github.com/nanobox-io/nanobox/util"
	"github.com/nanobox-io/nanobox/util/display"
	"github.com/nanobox-io/nanobox/util/provider"
	//	"github.com/nanobox-io/nanobox/util/update"
)

func Update() error {

	// init docker client
	if err := process_provider.Init(); err != nil {
		return fmt.Errorf("failed to init docker client: %s", err.Error())
	}

	// // check to see if nanobox needs to update
	// update.Check()

	// update all the nanobox images
	if err := pullImages(); err != nil {
		return fmt.Errorf("failed to pull images: %s", err)
	}

	// pull the latest docker-machine image
	return provider.Install()
}

func pullImages() error {
	display.OpenContext("Updating Images")
	defer display.CloseContext()

	images, err := docker.ImageList()
	if err != nil {
		return err
	}

	for _, image := range images {
		display.StartTask("Pulling %s image", image.Slug)

		// generate a docker percent display
		dockerPercent := &display.DockerPercentDisplay{
			Output: display.NewStreamer("info"),
			// Prefix: buildImage,
		}

		// pull the build image
		imagePullFunc := func() error {
			_, err := docker.ImagePull(image.Slug, dockerPercent)
			return err
		}

		if err := util.Retry(imagePullFunc, 5, time.Second); err != nil {
			lumber.Error("code:pullBuildImage:docker.ImagePull(%s, nil): %s", image.Slug, err.Error())
			display.ErrorTask()
			return fmt.Errorf("failed to pull docker image (%s): %s", image.Slug, err.Error())
		}

		display.StopTask()
	}

	return nil
}
<|endoftext|>


<commit_before>package database

import (
	"fmt"
	"kmud/utils"
	"labix.org/v2/mgo/bson"
	"strings"
	"sync"
)

type Coordinate struct {
	X int
	Y int
	Z int
}

// All database types should meet this interface
type Identifiable interface {
	GetId() bson.ObjectId
	GetType() objectType
}

type Nameable interface {
	PrettyName() string
}

type objectType int

const (
	characterType objectType = iota
	roomType      objectType = iota
)

const (
	dbObjectName string = "name"
)

type DbObject struct {
	Id      bson.ObjectId `bson:"_id"`
	objType objectType
	Name    string `bson:",omitempty"`
	Fields  map[string]interface{}
	mutex   sync.Mutex
}

type User struct {
	DbObject  `bson:",inline"`
	ColorMode utils.ColorMode
	online    bool
}

func NewUser(name string) User {
	var user User
	user.Id = bson.NewObjectId()
	user.Name = name
	user.ColorMode = utils.ColorModeNone
	return user
}

type ObjectField string

type Zone struct {
	DbObject `bson:",inline"`
}

func NewZone(name string) Zone {
	var zone Zone
	zone.Id = bson.NewObjectId()
	zone.Name = name
	return zone
}

type Item struct {
	DbObject `bson:",inline"`
}

func NewItem(name string) Item {
	var item Item
	item.Id = bson.NewObjectId()
	item.Name = name

	return item
}

type ExitDirection int

const (
	DirectionNone      ExitDirection = iota
	DirectionNorth     ExitDirection = iota
	DirectionNorthEast ExitDirection = iota
	DirectionEast      ExitDirection = iota
	DirectionSouthEast ExitDirection = iota
	DirectionSouth     ExitDirection = iota
	DirectionSouthWest ExitDirection = iota
	DirectionWest      ExitDirection = iota
	DirectionNorthWest ExitDirection = iota
	DirectionUp        ExitDirection = iota
	DirectionDown      ExitDirection = iota
)

type PrintMode int

const (
	ReadMode PrintMode = iota
	EditMode PrintMode = iota
)

func directionToExitString(colorMode utils.ColorMode, direction ExitDirection) string {

	letterColor := utils.ColorBlue
	bracketColor := utils.ColorDarkBlue
	textColor := utils.ColorWhite

	colorize := func(letters string, text string) string {
		return fmt.Sprintf("%s%s%s%s",
			utils.Colorize(colorMode, bracketColor, "["),
			utils.Colorize(colorMode, letterColor, letters),
			utils.Colorize(colorMode, bracketColor, "]"),
			utils.Colorize(colorMode, textColor, text))
	}

	switch direction {
	case DirectionNorth:
		return colorize("N", "orth")
	case DirectionNorthEast:
		return colorize("NE", "North East")
	case DirectionEast:
		return colorize("E", "ast")
	case DirectionSouthEast:
		return colorize("SE", "South East")
	case DirectionSouth:
		return colorize("S", "outh")
	case DirectionSouthWest:
		return colorize("SW", "South West")
	case DirectionWest:
		return colorize("W", "est")
	case DirectionNorthWest:
		return colorize("NW", "North West")
	case DirectionUp:
		return colorize("U", "p")
	case DirectionDown:
		return colorize("D", "own")
	case DirectionNone:
		return utils.Colorize(colorMode, utils.ColorWhite, "None")
	}

	panic("Unexpected code path")
}

func (self *DbObject) initDbObject(objType objectType) {
	self.Id = bson.NewObjectId()
	self.objType = objType
	self.Fields = map[string]interface{}{}
}

func (self DbObject) GetId() bson.ObjectId {
	return self.Id
}

func (self DbObject) GetType() objectType {
	return self.objType
}

func (self DbObject) PrettyName() string {
	return utils.FormatName(self.Name)
}

func (self *DbObject) setField(key string, value interface{}) {
	self.mutex.Lock()
	defer self.mutex.Unlock()

	self.Fields[key] = value
	updateObject(*self, "fields."+key, value)
}

func (self *DbObject) getField(key string) interface{} {
	return self.Fields[key]
}

func (self *DbObject) hasField(key string) bool {
	_, found := self.Fields[key]
	return found
}

func (self *DbObject) SetName(name string) {
	self.Name = name
	updateObject(*self, dbObjectName, name)
}

func (self *User) SetOnline(online bool) {
	self.online = online
}

func (self *User) Online() bool {
	return self.online
}

func (self *Coordinate) Next(direction ExitDirection) Coordinate {
	newCoord := *self
	switch direction {
	case DirectionNorth:
		newCoord.Y -= 1
	case DirectionNorthEast:
		newCoord.Y -= 1
		newCoord.X += 1
	case DirectionEast:
		newCoord.X += 1
	case DirectionSouthEast:
		newCoord.Y += 1
		newCoord.X += 1
	case DirectionSouth:
		newCoord.Y += 1
	case DirectionSouthWest:
		newCoord.Y += 1
		newCoord.X -= 1
	case DirectionWest:
		newCoord.X -= 1
	case DirectionNorthWest:
		newCoord.Y -= 1
		newCoord.X -= 1
	case DirectionUp:
		newCoord.Z -= 1
	case DirectionDown:
		newCoord.Z += 1
	}
	return newCoord
}

func StringToDirection(str string) ExitDirection {
	dirStr := strings.ToLower(str)
	switch dirStr {
	case "n":
		fallthrough
	case "north":
		return DirectionNorth
	case "ne":
		return DirectionNorthEast
	case "e":
		fallthrough
	case "east":
		return DirectionEast
	case "se":
		return DirectionSouthEast
	case "s":
		fallthrough
	case "south":
		return DirectionSouth
	case "sw":
		return DirectionSouthWest
	case "w":
		fallthrough
	case "west":
		return DirectionWest
	case "nw":
		return DirectionNorthWest
	case "u":
		fallthrough
	case "up":
		return DirectionUp
	case "d":
		fallthrough
	case "down":
		return DirectionDown
	}

	return DirectionNone
}

// vim: nocindent
<commit_msg>Even more thread safety<commit_after>package database

import (
	"fmt"
	"kmud/utils"
	"labix.org/v2/mgo/bson"
	"strings"
	"sync"
)

type Coordinate struct {
	X int
	Y int
	Z int
}

// All database types should meet this interface
type Identifiable interface {
	GetId() bson.ObjectId
	GetType() objectType
}

type Nameable interface {
	PrettyName() string
}

type objectType int

const (
	characterType objectType = iota
	roomType      objectType = iota
)

const (
	dbObjectName string = "name"
)

type DbObject struct {
	Id      bson.ObjectId `bson:"_id"`
	objType objectType
	Name    string `bson:",omitempty"`
	Fields  map[string]interface{}
	mutex   sync.Mutex
}

type User struct {
	DbObject  `bson:",inline"`
	ColorMode utils.ColorMode
	online    bool
}

func NewUser(name string) User {
	var user User
	user.Id = bson.NewObjectId()
	user.Name = name
	user.ColorMode = utils.ColorModeNone
	return user
}

type ObjectField string

type Zone struct {
	DbObject `bson:",inline"`
}

func NewZone(name string) Zone {
	var zone Zone
	zone.Id = bson.NewObjectId()
	zone.Name = name
	return zone
}

type Item struct {
	DbObject `bson:",inline"`
}

func NewItem(name string) Item {
	var item Item
	item.Id = bson.NewObjectId()
	item.Name = name

	return item
}

type ExitDirection int

const (
	DirectionNone      ExitDirection = iota
	DirectionNorth     ExitDirection = iota
	DirectionNorthEast ExitDirection = iota
	DirectionEast      ExitDirection = iota
	DirectionSouthEast ExitDirection = iota
	DirectionSouth     ExitDirection = iota
	DirectionSouthWest ExitDirection = iota
	DirectionWest      ExitDirection = iota
	DirectionNorthWest ExitDirection = iota
	DirectionUp        ExitDirection = iota
	DirectionDown      ExitDirection = iota
)

type PrintMode int

const (
	ReadMode PrintMode = iota
	EditMode PrintMode = iota
)

func directionToExitString(colorMode utils.ColorMode, direction ExitDirection) string {

	letterColor := utils.ColorBlue
	bracketColor := utils.ColorDarkBlue
	textColor := utils.ColorWhite

	colorize := func(letters string, text string) string {
		return fmt.Sprintf("%s%s%s%s",
			utils.Colorize(colorMode, bracketColor, "["),
			utils.Colorize(colorMode, letterColor, letters),
			utils.Colorize(colorMode, bracketColor, "]"),
			utils.Colorize(colorMode, textColor, text))
	}

	switch direction {
	case DirectionNorth:
		return colorize("N", "orth")
	case DirectionNorthEast:
		return colorize("NE", "North East")
	case DirectionEast:
		return colorize("E", "ast")
	case DirectionSouthEast:
		return colorize("SE", "South East")
	case DirectionSouth:
		return colorize("S", "outh")
	case DirectionSouthWest:
		return colorize("SW", "South West")
	case DirectionWest:
		return colorize("W", "est")
	case DirectionNorthWest:
		return colorize("NW", "North West")
	case DirectionUp:
		return colorize("U", "p")
	case DirectionDown:
		return colorize("D", "own")
	case DirectionNone:
		return utils.Colorize(colorMode, utils.ColorWhite, "None")
	}

	panic("Unexpected code path")
}

func (self *DbObject) initDbObject(objType objectType) {
	self.Id = bson.NewObjectId()
	self.objType = objType
	self.Fields = map[string]interface{}{}
}

func (self DbObject) GetId() bson.ObjectId {
	return self.Id
}

func (self DbObject) GetType() objectType {
	return self.objType
}

func (self DbObject) PrettyName() string {
	return utils.FormatName(self.Name)
}

func (self *DbObject) setField(key string, value interface{}) {
	self.mutex.Lock()
	defer self.mutex.Unlock()

	self.Fields[key] = value
	updateObject(*self, "fields."+key, value)
}

func (self *DbObject) getField(key string) interface{} {
    self.mutex.Lock()
    defer self.mutex.Unlock()

	return self.Fields[key]
}

func (self *DbObject) hasField(key string) bool {
	_, found := self.Fields[key]
	return found
}

func (self *DbObject) SetName(name string) {
	self.Name = name
	updateObject(*self, dbObjectName, name)
}

func (self *User) SetOnline(online bool) {
	self.online = online
}

func (self *User) Online() bool {
	return self.online
}

func (self *Coordinate) Next(direction ExitDirection) Coordinate {
	newCoord := *self
	switch direction {
	case DirectionNorth:
		newCoord.Y -= 1
	case DirectionNorthEast:
		newCoord.Y -= 1
		newCoord.X += 1
	case DirectionEast:
		newCoord.X += 1
	case DirectionSouthEast:
		newCoord.Y += 1
		newCoord.X += 1
	case DirectionSouth:
		newCoord.Y += 1
	case DirectionSouthWest:
		newCoord.Y += 1
		newCoord.X -= 1
	case DirectionWest:
		newCoord.X -= 1
	case DirectionNorthWest:
		newCoord.Y -= 1
		newCoord.X -= 1
	case DirectionUp:
		newCoord.Z -= 1
	case DirectionDown:
		newCoord.Z += 1
	}
	return newCoord
}

func StringToDirection(str string) ExitDirection {
	dirStr := strings.ToLower(str)
	switch dirStr {
	case "n":
		fallthrough
	case "north":
		return DirectionNorth
	case "ne":
		return DirectionNorthEast
	case "e":
		fallthrough
	case "east":
		return DirectionEast
	case "se":
		return DirectionSouthEast
	case "s":
		fallthrough
	case "south":
		return DirectionSouth
	case "sw":
		return DirectionSouthWest
	case "w":
		fallthrough
	case "west":
		return DirectionWest
	case "nw":
		return DirectionNorthWest
	case "u":
		fallthrough
	case "up":
		return DirectionUp
	case "d":
		fallthrough
	case "down":
		return DirectionDown
	}

	return DirectionNone
}

// vim: nocindent
<|endoftext|>


<commit_before>package producer

import (
	"errors"
	"fmt"
	"time"

	"h12.me/kpax/model"
	"h12.me/kpax/proto"
)

var (
	ErrProduceFailed    = errors.New("produce failed")
	ErrNoValidPartition = errors.New("no valid partition")
)

type P struct {
	LeaderRecoveryTime time.Duration
	RequiredAcks       int16
	AckTimeout         time.Duration
	cluster            model.Cluster
	topicPartitioner   *topicPartitioner
}

func New(cluster model.Cluster) *P {
	return &P{
		cluster:            cluster,
		topicPartitioner:   newTopicPartitioner(),
		LeaderRecoveryTime: 60 * time.Second,
		RequiredAcks:       proto.AckLocal,
		AckTimeout:         10 * time.Second,
	}
}

func (p *P) Produce(topic string, key, value []byte) error {
	partitioner := p.topicPartitioner.Get(topic)
	if partitioner == nil {
		partitions, err := p.cluster.Partitions(topic)
		if err != nil {
			return err
		}
		partitioner = p.topicPartitioner.Add(topic, partitions)
	}
	messageSet := getMessageSet(key, value)
nextPartition:
	for i := 0; i < partitioner.Count(); i++ {
		partition, err := partitioner.Partition(key)
		if err != nil {
			p.topicPartitioner.Delete(topic)
			return err
		}
		if err := (&proto.Payload{
			Topic:        topic,
			Partition:    partition,
			MessageSet:   messageSet,
			RequiredAcks: p.RequiredAcks,
			AckTimeout:   p.AckTimeout,
		}).Produce(p.cluster); err != nil {
			partitioner.Skip(partition)
			continue nextPartition
		}
		return nil
	}
	return fmt.Errorf("fail to produce to all partitions in %s", topic)
}

func (p *P) ProduceWithPartition(topic string, partition int32, key, value []byte) error {
	messageSet := getMessageSet(key, value)
	return (&proto.Payload{
		Topic:        topic,
		Partition:    partition,
		MessageSet:   messageSet,
		RequiredAcks: p.RequiredAcks,
		AckTimeout:   p.AckTimeout,
	}).Produce(p.cluster)
}

func getMessageSet(key, value []byte) []proto.OffsetMessage {
	return []proto.OffsetMessage{
		{
			SizedMessage: proto.SizedMessage{CRCMessage: proto.CRCMessage{
				Message: proto.Message{
					Key:   key,
					Value: value,
				},
			}}},
	}
}
<commit_msg>support ProduceMessageSet method<commit_after>package producer

import (
	"errors"
	"fmt"
	"time"

	"h12.me/kpax/model"
	"h12.me/kpax/proto"
)

var (
	ErrProduceFailed    = errors.New("produce failed")
	ErrNoValidPartition = errors.New("no valid partition")
)

type P struct {
	LeaderRecoveryTime time.Duration
	RequiredAcks       proto.ProduceAckType
	AckTimeout         time.Duration
	cluster            model.Cluster
	topicPartitioner   *topicPartitioner
}

func New(cluster model.Cluster) *P {
	return &P{
		cluster:            cluster,
		topicPartitioner:   newTopicPartitioner(),
		LeaderRecoveryTime: 60 * time.Second,
		RequiredAcks:       proto.AckLocal,
		AckTimeout:         10 * time.Second,
	}
}

func (p *P) ProduceMessageSet(topic string, messageSet proto.MessageSet) error {
	if len(messageSet) == 0 {
		panic("empty message set")
	}
	key := messageSet[0].Key
	partitioner := p.topicPartitioner.Get(topic)
	if partitioner == nil {
		partitions, err := p.cluster.Partitions(topic)
		if err != nil {
			return err
		}
		partitioner = p.topicPartitioner.Add(topic, partitions)
	}
nextPartition:
	for i := 0; i < partitioner.Count(); i++ {
		partition, err := partitioner.Partition(key)
		if err != nil {
			p.topicPartitioner.Delete(topic)
			return err
		}
		if err := (&proto.Payload{
			Topic:        topic,
			Partition:    partition,
			MessageSet:   messageSet,
			RequiredAcks: p.RequiredAcks,
			AckTimeout:   p.AckTimeout,
		}).Produce(p.cluster); err != nil {
			partitioner.Skip(partition)
			continue nextPartition
		}
		return nil
	}
	return fmt.Errorf("fail to produce to all partitions in %s", topic)
}

func (p *P) Produce(topic string, key, value []byte) error {
	return p.ProduceMessageSet(topic, getMessageSet(key, value))
}

func (p *P) ProduceWithPartition(topic string, partition int32, key, value []byte) error {
	messageSet := getMessageSet(key, value)
	return (&proto.Payload{
		Topic:        topic,
		Partition:    partition,
		MessageSet:   messageSet,
		RequiredAcks: p.RequiredAcks,
		AckTimeout:   p.AckTimeout,
	}).Produce(p.cluster)
}

func getMessageSet(key, value []byte) []proto.OffsetMessage {
	return []proto.OffsetMessage{
		{
			SizedMessage: proto.SizedMessage{CRCMessage: proto.CRCMessage{
				Message: proto.Message{
					Key:   key,
					Value: value,
				},
			}}},
	}
}
<|endoftext|>