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

package alarm

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/robertkrimen/otto"
	"github.com/tsuru/tsuru-autoscale/action"
	"github.com/tsuru/tsuru-autoscale/datasource"
	"github.com/tsuru/tsuru-autoscale/db"
	"github.com/tsuru/tsuru-autoscale/log"
	"github.com/tsuru/tsuru-autoscale/tsuru"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func StartAutoScale() {
	runAutoScale()
}

func logger() *log.Logger {
	return log.Log()
}

// Alarm represents the configuration for the auto scale.
type Alarm struct {
	Name        string            `json:"name"`
	Actions     []string          `json:"actions"`
	Expression  string            `json:"expression"`
	Enabled     bool              `json:"enabled"`
	Wait        time.Duration     `json:"wait"`
	DataSources []string          `json:"datasources"`
	Instance    string            `json:"instance"`
	Envs        map[string]string `json:"envs"`
}

func NewAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Insert(&a)
}

func runAutoScaleOnce() {
	logger().Print("checking alarms")
	alarms := []Alarm{}
	conn, err := db.Conn()
	if err != nil {
		return
	}
	defer conn.Close()
	err = conn.Alarms().Find(bson.M{"enabled": true}).All(&alarms)
	if err != nil {
		logger().Error(err)
		return
	}
	var wg sync.WaitGroup
	for _, alarm := range alarms {
		wg.Add(1)
		go func(alarm Alarm) {
			defer wg.Done()
			logger().Printf("checking %s alarm", alarm.Name)
			err := scaleIfNeeded(&alarm)
			if err != nil {
				logger().Error(err)
			}
		}(alarm)
	}
	wg.Wait()
}

func interval() time.Duration {
	if i := os.Getenv("AUTOSCALE_INTERVAL"); i != "" {
		v, err := strconv.Atoi(i)
		if err == nil {
			return time.Duration(v)
		}
		logger().Error(err)
	}
	return time.Duration(10)
}

func runAutoScale() {
	for {
		runAutoScaleOnce()
		time.Sleep(interval() * time.Second)
	}
}

func scaleIfNeeded(alarm *Alarm) error {
	if alarm == nil {
		return errors.New("alarm: alarm is not configured.")
	}
	check, err := alarm.Check()
	if err != nil {
		logger().Error(err)
		return err
	}
	logger().Printf("alarm %s - %s - check: %t", alarm.Name, alarm.Expression, check)
	if check {
		if wait, err := shouldWait(alarm); err != nil {
			logger().Printf("waiting for alarm %s", alarm.Name)
			return err
		} else if wait {
			return nil
		}
		for _, alarmName := range alarm.Actions {
			a, err := action.FindByName(alarmName)
			if err != nil {
				logger().Error(err)
			} else {
				logger().Printf("executing alarm %s action %s", alarm.Name, a.Name)
				instance, err := tsuru.GetInstanceByName(alarm.Instance)
				if err != nil {
					logger().Error(err)
					return err
				}
				if len(instance.Apps) < 1 {
					msg := "Error trying to get app instance, auto scale aborted."
					logger().Printf(msg)
					err = errors.New(msg)
					return err
				}
				appName := instance.Apps[0]
				evt, err := NewEvent(alarm, a)
				if err != nil {
					logger().Error(err)
				}
				aErr := a.Do(appName, alarm.Envs)
				if aErr != nil {
					logger().Error(aErr)
				} else {
					logger().Printf("alarm %s action %s executed", alarm.Name, a.Name)
				}
				err = evt.update(aErr)
				if err != nil {
					logger().Error(err)
				}
			}
		}
		return nil
	}
	return nil
}

func shouldWait(alarm *Alarm) (bool, error) {
	now := time.Now().UTC()
	lastEvent, err := lastScaleEvent(alarm)
	if err != nil && err != mgo.ErrNotFound {
		logger().Error(err)
		return false, err
	}
	if err != mgo.ErrNotFound && lastEvent.EndTime.IsZero() {
		logger().Printf("last event not finished yet for alarm %s - waiting", alarm.Name)
		return true, nil
	}
	diff := now.Sub(lastEvent.EndTime)
	if diff > alarm.Wait {
		logger().Printf("diff %d > %d form alarm %s - not waiting", diff, alarm.Wait, alarm.Name)
		return false, nil
	}
	logger().Printf("diff %d < %d form alarm %s - waiting", diff, alarm.Wait, alarm.Name)
	return true, nil
}

func Enable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": true}})
}

func Disable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": false}})
}

func (a *Alarm) data(appName string) (map[string]string, error) {
	d := map[string]string{}
	for _, dataSource := range a.DataSources {
		ds, err := datasource.Get(dataSource)
		if err != nil {
			logger().Error(err)
			return nil, err
		}
		data, err := ds.Get(appName, a.Envs)
		if err != nil {
			logger().Error(err)
			return nil, err
		}
		logger().Printf("data for alarm %s - %s", a.Name, data)
		d[ds.Name] = data
	}
	return d, nil
}

func (a *Alarm) Check() (bool, error) {
	instance, err := tsuru.GetInstanceByName(a.Instance)
	if err != nil {
		logger().Error(err)
		return false, err
	}
	if len(instance.Apps) < 1 {
		msg := "Error trying to get app instance."
		logger().Printf(msg)
		err = errors.New(msg)
		return false, err
	}
	appName := instance.Apps[0]
	dataSourceData, err := a.data(appName)
	if err != nil {
		logger().Error(err)
		return false, err
	}
	expression := strings.Replace(a.Expression, "{app}", appName, -1)
	for key, value := range a.Envs {
		expression = strings.Replace(expression, fmt.Sprintf("{%s}", key), value, -1)
	}
	data := ""
	for key, value := range dataSourceData {
		data += fmt.Sprintf("var %s=%s;", key, value)
	}
	vm := otto.New()
	vm.Run(data)
	vm.Run(fmt.Sprintf("var expression=%s;", expression))
	result, err := vm.Get("expression")
	if err != nil {
		logger().Error(err)
		return false, err
	}
	check, err := result.ToBoolean()
	if err != nil {
		logger().Error(err)
		return false, err
	}
	return check, nil
}

// ListAlarmsByToken lists alarms by token.
func ListAlarmsByToken(token string) ([]Alarm, error) {
	i, err := tsuru.FindServiceInstance(token)
	if err != nil {
		return nil, err
	}
	instances := []string{}
	for _, instance := range i {
		instances = append(instances, instance.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": bson.M{"$in": instances}}).All(&alarms)
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	return alarms, nil
}

// ListAlarmsByInstance lists alarms by instance.
func ListAlarmsByInstance(instanceName string) ([]Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": instanceName}).All(&alarms)
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	return alarms, nil
}

// FindAlarmByName find alarm by name.
func FindAlarmByName(name string) (*Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	defer conn.Close()
	var alarm Alarm
	err = conn.Alarms().Find(bson.M{"name": name}).One(&alarm)
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	return &alarm, nil
}

// RemoveAlarm removes an alarm.
func RemoveAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return err
	}
	defer conn.Close()
	err = conn.Alarms().Remove(bson.M{"name": a.Name})
	if err != nil {
		logger().Error(err)
		return err
	}
	conn.Events().RemoveAll(bson.M{"alarm.name": a.Name})
	return nil
}
<commit_msg>add log on connection error<commit_after>// Copyright 2015 tsuru-autoscale authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package alarm

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/robertkrimen/otto"
	"github.com/tsuru/tsuru-autoscale/action"
	"github.com/tsuru/tsuru-autoscale/datasource"
	"github.com/tsuru/tsuru-autoscale/db"
	"github.com/tsuru/tsuru-autoscale/log"
	"github.com/tsuru/tsuru-autoscale/tsuru"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func StartAutoScale() {
	runAutoScale()
}

func logger() *log.Logger {
	return log.Log()
}

// Alarm represents the configuration for the auto scale.
type Alarm struct {
	Name        string            `json:"name"`
	Actions     []string          `json:"actions"`
	Expression  string            `json:"expression"`
	Enabled     bool              `json:"enabled"`
	Wait        time.Duration     `json:"wait"`
	DataSources []string          `json:"datasources"`
	Instance    string            `json:"instance"`
	Envs        map[string]string `json:"envs"`
}

func NewAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Insert(&a)
}

func runAutoScaleOnce() {
	logger().Print("checking alarms")
	alarms := []Alarm{}
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return
	}
	defer conn.Close()
	err = conn.Alarms().Find(bson.M{"enabled": true}).All(&alarms)
	if err != nil {
		logger().Error(err)
		return
	}
	var wg sync.WaitGroup
	for _, alarm := range alarms {
		wg.Add(1)
		go func(alarm Alarm) {
			defer wg.Done()
			logger().Printf("checking %s alarm", alarm.Name)
			err := scaleIfNeeded(&alarm)
			if err != nil {
				logger().Error(err)
			}
		}(alarm)
	}
	wg.Wait()
}

func interval() time.Duration {
	if i := os.Getenv("AUTOSCALE_INTERVAL"); i != "" {
		v, err := strconv.Atoi(i)
		if err == nil {
			return time.Duration(v)
		}
		logger().Error(err)
	}
	return time.Duration(10)
}

func runAutoScale() {
	for {
		runAutoScaleOnce()
		time.Sleep(interval() * time.Second)
	}
}

func scaleIfNeeded(alarm *Alarm) error {
	if alarm == nil {
		return errors.New("alarm: alarm is not configured.")
	}
	check, err := alarm.Check()
	if err != nil {
		logger().Error(err)
		return err
	}
	logger().Printf("alarm %s - %s - check: %t", alarm.Name, alarm.Expression, check)
	if check {
		if wait, err := shouldWait(alarm); err != nil {
			logger().Printf("waiting for alarm %s", alarm.Name)
			return err
		} else if wait {
			return nil
		}
		for _, alarmName := range alarm.Actions {
			a, err := action.FindByName(alarmName)
			if err != nil {
				logger().Error(err)
			} else {
				logger().Printf("executing alarm %s action %s", alarm.Name, a.Name)
				instance, err := tsuru.GetInstanceByName(alarm.Instance)
				if err != nil {
					logger().Error(err)
					return err
				}
				if len(instance.Apps) < 1 {
					msg := "Error trying to get app instance, auto scale aborted."
					logger().Printf(msg)
					err = errors.New(msg)
					return err
				}
				appName := instance.Apps[0]
				evt, err := NewEvent(alarm, a)
				if err != nil {
					logger().Error(err)
				}
				aErr := a.Do(appName, alarm.Envs)
				if aErr != nil {
					logger().Error(aErr)
				} else {
					logger().Printf("alarm %s action %s executed", alarm.Name, a.Name)
				}
				err = evt.update(aErr)
				if err != nil {
					logger().Error(err)
				}
			}
		}
		return nil
	}
	return nil
}

func shouldWait(alarm *Alarm) (bool, error) {
	now := time.Now().UTC()
	lastEvent, err := lastScaleEvent(alarm)
	if err != nil && err != mgo.ErrNotFound {
		logger().Error(err)
		return false, err
	}
	if err != mgo.ErrNotFound && lastEvent.EndTime.IsZero() {
		logger().Printf("last event not finished yet for alarm %s - waiting", alarm.Name)
		return true, nil
	}
	diff := now.Sub(lastEvent.EndTime)
	if diff > alarm.Wait {
		logger().Printf("diff %d > %d form alarm %s - not waiting", diff, alarm.Wait, alarm.Name)
		return false, nil
	}
	logger().Printf("diff %d < %d form alarm %s - waiting", diff, alarm.Wait, alarm.Name)
	return true, nil
}

func Enable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": true}})
}

func Disable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": false}})
}

func (a *Alarm) data(appName string) (map[string]string, error) {
	d := map[string]string{}
	for _, dataSource := range a.DataSources {
		ds, err := datasource.Get(dataSource)
		if err != nil {
			logger().Error(err)
			return nil, err
		}
		data, err := ds.Get(appName, a.Envs)
		if err != nil {
			logger().Error(err)
			return nil, err
		}
		logger().Printf("data for alarm %s - %s", a.Name, data)
		d[ds.Name] = data
	}
	return d, nil
}

func (a *Alarm) Check() (bool, error) {
	instance, err := tsuru.GetInstanceByName(a.Instance)
	if err != nil {
		logger().Error(err)
		return false, err
	}
	if len(instance.Apps) < 1 {
		msg := "Error trying to get app instance."
		logger().Printf(msg)
		err = errors.New(msg)
		return false, err
	}
	appName := instance.Apps[0]
	dataSourceData, err := a.data(appName)
	if err != nil {
		logger().Error(err)
		return false, err
	}
	expression := strings.Replace(a.Expression, "{app}", appName, -1)
	for key, value := range a.Envs {
		expression = strings.Replace(expression, fmt.Sprintf("{%s}", key), value, -1)
	}
	data := ""
	for key, value := range dataSourceData {
		data += fmt.Sprintf("var %s=%s;", key, value)
	}
	vm := otto.New()
	vm.Run(data)
	vm.Run(fmt.Sprintf("var expression=%s;", expression))
	result, err := vm.Get("expression")
	if err != nil {
		logger().Error(err)
		return false, err
	}
	check, err := result.ToBoolean()
	if err != nil {
		logger().Error(err)
		return false, err
	}
	return check, nil
}

// ListAlarmsByToken lists alarms by token.
func ListAlarmsByToken(token string) ([]Alarm, error) {
	i, err := tsuru.FindServiceInstance(token)
	if err != nil {
		return nil, err
	}
	instances := []string{}
	for _, instance := range i {
		instances = append(instances, instance.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": bson.M{"$in": instances}}).All(&alarms)
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	return alarms, nil
}

// ListAlarmsByInstance lists alarms by instance.
func ListAlarmsByInstance(instanceName string) ([]Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": instanceName}).All(&alarms)
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	return alarms, nil
}

// FindAlarmByName find alarm by name.
func FindAlarmByName(name string) (*Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	defer conn.Close()
	var alarm Alarm
	err = conn.Alarms().Find(bson.M{"name": name}).One(&alarm)
	if err != nil {
		logger().Error(err)
		return nil, err
	}
	return &alarm, nil
}

// RemoveAlarm removes an alarm.
func RemoveAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		logger().Error(err)
		return err
	}
	defer conn.Close()
	err = conn.Alarms().Remove(bson.M{"name": a.Name})
	if err != nil {
		logger().Error(err)
		return err
	}
	conn.Events().RemoveAll(bson.M{"alarm.name": a.Name})
	return nil
}
<|endoftext|>


<commit_before>package bot

import (
	"math/rand"
	"regexp"
	"strings"
	"time"

	"fmt"

	. "github.com/PoolC/slack_bot/util"
	slack "github.com/nlopes/slack"
)

var (
	remember_re   *regexp.Regexp = regexp.MustCompile("^안즈쨩? 기억해? ([^/]+)/(.+)")
	tell_re       *regexp.Regexp = regexp.MustCompile("^안즈쨩? 알려줘 (.+)")
	kawaii_re     *regexp.Regexp = regexp.MustCompile("^안즈쨩? 카와이")
	give_candy_re *regexp.Regexp = regexp.MustCompile("^안즈쨩? 사탕줄게")
)

type Anzu struct {
	*BaseBot
	rc RedisClient
}

func NewAnzu(token string, stop *chan struct{}, redisClient RedisClient) *Anzu {
	return &Anzu{NewBot(token, stop), redisClient}
}

func anzuMessageProcess(bot *Anzu, e *slack.MessageEvent) interface{} {
	force_accept := false
	switch {
	case e.Text == "사람은 일을 하고 살아야한다. 메우":
		return "이거 놔라 이 퇴근도 못하는 놈이"
	case e.Text == "안즈쨩 뭐해?":
		return "숨셔"
	default:
		if AcceptRE(e.Text, give_candy_re) {
			last := bot.rc.Get(fmt.Sprintf("%s_lastfail", e.User)).Val()
			if last == "" {
				return ""
			}
			force_accept = true
			e.Text = last
		}

		if matched, ok := MatchRE(e.Text, remember_re); ok {
			key, val := strings.TrimSpace(matched[1]), strings.TrimSpace(matched[2])
			var ret string
			switch {
			case key == "" || val == "":
				ret = "에...?"
			case AcceptRE(val, tell_re):
				ret = "에... 귀찮아..."
			case force_accept:
				ret = "응응 기억했어"
				fallthrough
			case rand.Float32() < 0.4:
				bot.rc.Set(key, val, 0)
				if len(ret) == 0 {
					ret = "에... 귀찮지만 기억했어"
				}
			default:
				ret = "귀찮아..."
				bot.rc.Set(fmt.Sprintf("%s_lastfail", e.User), e.Text, time.Duration(300*time.Second))
			}

			return ret
		} else if matched, ok := MatchRE(e.Text, tell_re); ok {
			key := strings.TrimSpace(matched[1])
			val := bot.rc.Get(key).Val()
			var ret string
			switch {
			case val == "":
				ret = "그런거 몰라"
			case force_accept:
				ret = fmt.Sprintf("%s 물어봤지? %s 야", key, val)
			case rand.Float32() < 0.4:
				ret = val
			default:
				bot.rc.Set(fmt.Sprintf("%s_lastfail", e.User), e.Text, time.Duration(300*time.Second))
				ret = "Zzz..."
			}
			return ret
		} else if _, ok := MatchRE(e.Text, kawaii_re); ok {
			return "뭐... 뭐라는거야"
		}
	}
	return nil
}

func (bot *Anzu) onMessageEvent(e *slack.MessageEvent) {
	message := anzuMessageProcess(bot, e)
	switch message.(type) {
	case string:
		bot.sendSimple(e, message.(string))
	}
}
<commit_msg>dice 기능 추가<commit_after>package bot

import (
	"math/rand"
	"regexp"
	"strings"
	"time"

	"fmt"

	. "github.com/PoolC/slack_bot/util"
	slack "github.com/nlopes/slack"
)

var (
	remember_re   *regexp.Regexp = regexp.MustCompile("^안즈쨩? 기억해? ([^/]+)/(.+)")
	tell_re       *regexp.Regexp = regexp.MustCompile("^안즈쨩? 알려줘 (.+)")
	kawaii_re     *regexp.Regexp = regexp.MustCompile("^안즈쨩? 카와이")
	give_candy_re *regexp.Regexp = regexp.MustCompile("^안즈쨩? 사탕줄게")
	dice          *regexp.Regexp = regexp.MustCompile("^안즈쨩? 주사위 ([0-9]+)")
)

type Anzu struct {
	*BaseBot
	rc RedisClient
}

func NewAnzu(token string, stop *chan struct{}, redisClient RedisClient) *Anzu {
	return &Anzu{NewBot(token, stop), redisClient}
}

func anzuMessageProcess(bot *Anzu, e *slack.MessageEvent) interface{} {
	force_accept := false
	switch {
	case e.Text == "사람은 일을 하고 살아야한다. 메우":
		return "이거 놔라 이 퇴근도 못하는 놈이"
	case e.Text == "안즈쨩 뭐해?":
		return "숨셔"
	default:
		if AcceptRE(e.Text, give_candy_re) {
			last := bot.rc.Get(fmt.Sprintf("%s_lastfail", e.User)).Val()
			if last == "" {
				return ""
			}
			force_accept = true
			e.Text = last
		}

		if matched, ok := MatchRE(e.Text, remember_re); ok {
			key, val := strings.TrimSpace(matched[1]), strings.TrimSpace(matched[2])
			var ret string
			switch {
			case key == "" || val == "":
				ret = "에...?"
			case AcceptRE(val, tell_re):
				ret = "에... 귀찮아..."
			case force_accept:
				ret = "응응 기억했어"
				fallthrough
			case rand.Float32() < 0.4:
				bot.rc.Set(key, val, 0)
				if len(ret) == 0 {
					ret = "에... 귀찮지만 기억했어"
				}
			default:
				ret = "귀찮아..."
				bot.rc.Set(fmt.Sprintf("%s_lastfail", e.User), e.Text, time.Duration(300*time.Second))
			}

			return ret
		} else if matched, ok := MatchRE(e.Text, tell_re); ok {
			key := strings.TrimSpace(matched[1])
			val := bot.rc.Get(key).Val()
			var ret string
			switch {
			case val == "":
				ret = "그런거 몰라"
			case force_accept:
				ret = fmt.Sprintf("%s 물어봤지? %s 야", key, val)
			case rand.Float32() < 0.4:
				ret = val
			default:
				bot.rc.Set(fmt.Sprintf("%s_lastfail", e.User), e.Text, time.Duration(300*time.Second))
				ret = "Zzz..."
			}
			return ret
		} else if _, ok := MatchRE(e.Text, kawaii_re); ok {
			return "뭐... 뭐라는거야"
		} else if matched, ok := MatchRE(e.Text, dice); ok {
			num := matched[1]
			return fmt.Sprintf("%d", num * rand.Float32())
		}
	}
	return nil
}

func (bot *Anzu) onMessageEvent(e *slack.MessageEvent) {
	message := anzuMessageProcess(bot, e)
	switch message.(type) {
	case string:
		bot.sendSimple(e, message.(string))
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Helton Marques
//
//	Use of this source code is governed by a LGPL
//	license that can be found in the LICENSE file.
//

package ami

import (
	"errors"
	"fmt"
	"os"
	"strings"
)

func VersionInfo() string {
	return "Version: 0.0.1; Build Date: Jan 14, 2013;"
}

func sendCmd(socket *Socket, cmd []string) error {
	for _, s := range cmd {
		err := socket.Send("%s", s)
		if err != nil {
			fmt.Printf("send login error:[%v]\n", err)
			return err
		}
	}
	return nil
}

func Connect(socket *Socket) (bool, error) {
	if !socket.Connected() {
		fmt.Printf("could not connect to AMI\n")
	}
	if answer, err := socket.Recv(); err != nil || !strings.Contains(answer, "Asterisk Call Manager") {
		return false, errors.New("AMI connection failed")
	}
	return true, nil
}

//  Login
//      Login Manager.
//
func Login(socket *Socket, user, secret, events, actionID string) (bool, error) {
	// verify parameters
	if len(user) == 0 || len(secret) == 0 || len(events) == 0 || len(actionID) == 0 {
		return false, errors.New("Invalid parameters")
	}

	command := []string{
		"Action: Login",
		"\r\nUsername: ",
		user,
		"\r\nSecret: ",
		secret,
		"\r\nEvents: ",
		events,
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	message, err := getMessage(socket, command, actionID)
	if err != nil {
		return false, err
	}

	if message["Response"] != "Success" {
		return false, errors.New(message["Message"])
	}
	return true, nil
}

//  Logoff
//      Logoff the current manager session.
//
func Logoff(socket *Socket, actionID string) (bool, error) {
	// verify parameters
	if len(actionID) == 0 {
		return false, errors.New("Invalid parameters")
	}

	command := []string{
		"Action: Logoff",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	message, err := getMessage(socket, command, actionID)
	if err != nil {
		return false, err
	}
	if message["Response"] != "Goodbye" {
		return false, errors.New(message["Message"])
	}
	return true, nil
}

func GetUUID() (string, error) {
	f, err := os.Open("/dev/urandom")
	if err != nil {
		fmt.Printf("open /dev/urandom error:[%v]\n", err)
		return "", err
	}
	b := make([]byte, 16)
	f.Read(b)
	f.Close()
	uuid := fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6],
		b[6:8], b[8:10], b[10:])
	return uuid, nil
}

//  Ping
//      A 'Ping' action will ellicit a 'Pong' response.
//      Used to keep the manager connection open.
//
func Ping(socket *Socket, actionID string) (bool, error) {

	// verify parameters
	if len(actionID) == 0 {
		return false, errors.New("Invalid parameters")
	}

	command := []string{
		"Action: Ping",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	message, err := getMessage(socket, command, actionID)
	if err != nil {
		return false, err
	}
	if message["Response"] != "Success" {
		return false, errors.New(message["Message"])
	}
	return true, nil
}

//	Challenge
//		Generate a challenge for MD5 authentication.
//
func Challenge(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Challenge",
		"\r\nActionID: ",
		actionID,
		"\r\nAuthType: ",
		"MD5",
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	Command
//		Execute Asterisk CLI Command.
//
func Command(socket *Socket, actionID, cmd string) (map[string]string, error) {
	if len(actionID) == 0 || len(cmd) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Command",
		"\r\nActionID: ",
		actionID,
		"\r\nCommand: ",
		cmd,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	CoreSettings
//		Show PBX core settings (version etc).
//
func CoreSettings(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: CoreSettings",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	CoreStatus
//		Show PBX core status variables.
//
func CoreStatus(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: CoreStatus",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	CreateConfig
//		Creates an empty file in the configuration directory.
//		This action will create an empty file in the configuration directory.
//		This action is intended to be used before an UpdateConfig action.
//
func CreateConfig(socket *Socket, actionID, filename string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: CreateConfig",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	DataGet
//		Retrieve the data api tree.
//
func DataGet(socket *Socket, actionID, path, search, filter string) (map[string]string, error) {
	if len(actionID) == 0 || len(path) == 0 ||
		len(search) == 0 || len(filter) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: DataGet",
		"\r\nActionID: ",
		actionID,
		"\r\nPath: ",
		path,
		"\r\nSearch: ",
		search,
		"\r\nFilter: ",
		filter,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	Events
//		Control Event Flow.
//		Enable/Disable sending of events to this manager client.
//
func Events(socket *Socket, actionID, eventMask string) (map[string]string, error) {
	if len(actionID) == 0 || len(eventMask) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Events",
		"\r\nActionID: ",
		actionID,
		"\r\nEventMask: ",
		eventMask,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	GetConfig
//		Retrieve configuration.
//		This action will dump the contents of a configuration file by category and contents or optionally by specified category only.
//
func GetConfig(socket *Socket, actionID, filename, category string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 || len(category) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: GetConfig",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\nCategory: ",
		category,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	GetConfigJSON
//		Retrieve configuration (JSON format).
//		This action will dump the contents of a configuration file by category and contents in JSON format.
//		This only makes sense to be used using rawman over the HTTP interface.
//
func GetConfigJSON(socket *Socket, actionID, filename string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: GetConfigJSON",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	JabberSend
//		Sends a message to a Jabber Client.
//
func JabberSend(socket *Socket, actionID, jabber, jid, message string) (map[string]string, error) {
	if len(actionID) == 0 || len(jabber) == 0 ||
		len(jid) == 0 || len(message) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: JabberSend",
		"\r\nActionID: ",
		actionID,
		"\r\nJabber: ",
		jabber,
		"\r\nJID: ",
		jid,
		"\r\nMessage: ",
		message,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ListCommands
//		List available manager commands.
//		Returns the action name and synopsis for every action that is available to the user
//
func ListCommands(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ListCommands",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ListCategories
//		List categories in configuration file.
//
func ListCategories(socket *Socket, actionID, filename string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ListCategories",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ModuleCheck
//		Check if module is loaded.
//		Checks if Asterisk module is loaded. Will return Success/Failure. For success returns, the module revision number is included.
//
func ModuleCheck(socket *Socket, actionID, module string) (map[string]string, error) {
	if len(actionID) == 0 || len(module) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ModuleCheck",
		"\r\nActionID: ",
		actionID,
		"\r\nModule: ",
		module,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ModuleLoad
//		Module management.
//		Loads, unloads or reloads an Asterisk module in a running system.
//
func ModuleLoad(socket *Socket, actionID, module, loadType string) (map[string]string, error) {
	if len(actionID) == 0 || len(module) == 0 || len(loadType) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ModuleLoad",
		"\r\nActionID: ",
		actionID,
		"\r\nModule: ",
		module,
		"\r\nLoadType: ",
		loadType,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	Reload
//		Send a reload event.
//
func Reload(socket *Socket, actionID, module string) (map[string]string, error) {
	if len(actionID) == 0 || len(module) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Reload",
		"\r\nActionID: ",
		actionID,
		"\r\nModule: ",
		module,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ShowDialPlan
//		Show dialplan contexts and extensions
//		Be aware that showing the full dialplan may take a lot of capacity
//
func ShowDialPlan(socket *Socket, actionID, extension, context string) (map[string]string, error) {
	if len(actionID) == 0 || len(extension) == 0 || len(context) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ShowDialPlan",
		"\r\nActionID: ",
		actionID,
		"\r\nExtension: ",
		extension,
		"\r\nContext: ",
		context,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	GetEvents
//		Get events from current socket connection
//		It is mandatory set 'events' of ami.Login with "system,call,all,user", to received
//		events
//
func GetEvents(socket *Socket) (map[string]string, error) {
	message, err := decode(socket)
	if err != nil {
		return nil, err
	}
	return message, nil
}
<commit_msg>Update manager.go<commit_after>// Copyright 2014 Helton Marques
//
//	Use of this source code is governed by a LGPL
//	license that can be found in the LICENSE file.
//

package ami

import (
	"errors"
	"fmt"
	"os"
	"strings"
)

func VersionInfo() string {
	return "Version: 1.0.0; Build Date: Jan 16, 2014;"
}

func sendCmd(socket *Socket, cmd []string) error {
	for _, s := range cmd {
		err := socket.Send("%s", s)
		if err != nil {
			fmt.Printf("send login error:[%v]\n", err)
			return err
		}
	}
	return nil
}

func Connect(socket *Socket) (bool, error) {
	if !socket.Connected() {
		fmt.Printf("could not connect to AMI\n")
	}
	if answer, err := socket.Recv(); err != nil || !strings.Contains(answer, "Asterisk Call Manager") {
		return false, errors.New("AMI connection failed")
	}
	return true, nil
}

//  Login
//      Login Manager.
//
func Login(socket *Socket, user, secret, events, actionID string) (bool, error) {
	// verify parameters
	if len(user) == 0 || len(secret) == 0 || len(events) == 0 || len(actionID) == 0 {
		return false, errors.New("Invalid parameters")
	}

	command := []string{
		"Action: Login",
		"\r\nUsername: ",
		user,
		"\r\nSecret: ",
		secret,
		"\r\nEvents: ",
		events,
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	message, err := getMessage(socket, command, actionID)
	if err != nil {
		return false, err
	}

	if message["Response"] != "Success" {
		return false, errors.New(message["Message"])
	}
	return true, nil
}

//  Logoff
//      Logoff the current manager session.
//
func Logoff(socket *Socket, actionID string) (bool, error) {
	// verify parameters
	if len(actionID) == 0 {
		return false, errors.New("Invalid parameters")
	}

	command := []string{
		"Action: Logoff",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	message, err := getMessage(socket, command, actionID)
	if err != nil {
		return false, err
	}
	if message["Response"] != "Goodbye" {
		return false, errors.New(message["Message"])
	}
	return true, nil
}

func GetUUID() (string, error) {
	f, err := os.Open("/dev/urandom")
	if err != nil {
		fmt.Printf("open /dev/urandom error:[%v]\n", err)
		return "", err
	}
	b := make([]byte, 16)
	f.Read(b)
	f.Close()
	uuid := fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6],
		b[6:8], b[8:10], b[10:])
	return uuid, nil
}

//  Ping
//      A 'Ping' action will ellicit a 'Pong' response.
//      Used to keep the manager connection open.
//
func Ping(socket *Socket, actionID string) (bool, error) {

	// verify parameters
	if len(actionID) == 0 {
		return false, errors.New("Invalid parameters")
	}

	command := []string{
		"Action: Ping",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	message, err := getMessage(socket, command, actionID)
	if err != nil {
		return false, err
	}
	if message["Response"] != "Success" {
		return false, errors.New(message["Message"])
	}
	return true, nil
}

//	Challenge
//		Generate a challenge for MD5 authentication.
//
func Challenge(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Challenge",
		"\r\nActionID: ",
		actionID,
		"\r\nAuthType: ",
		"MD5",
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	Command
//		Execute Asterisk CLI Command.
//
func Command(socket *Socket, actionID, cmd string) (map[string]string, error) {
	if len(actionID) == 0 || len(cmd) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Command",
		"\r\nActionID: ",
		actionID,
		"\r\nCommand: ",
		cmd,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	CoreSettings
//		Show PBX core settings (version etc).
//
func CoreSettings(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: CoreSettings",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	CoreStatus
//		Show PBX core status variables.
//
func CoreStatus(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: CoreStatus",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	CreateConfig
//		Creates an empty file in the configuration directory.
//		This action will create an empty file in the configuration directory.
//		This action is intended to be used before an UpdateConfig action.
//
func CreateConfig(socket *Socket, actionID, filename string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: CreateConfig",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	DataGet
//		Retrieve the data api tree.
//
func DataGet(socket *Socket, actionID, path, search, filter string) (map[string]string, error) {
	if len(actionID) == 0 || len(path) == 0 ||
		len(search) == 0 || len(filter) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: DataGet",
		"\r\nActionID: ",
		actionID,
		"\r\nPath: ",
		path,
		"\r\nSearch: ",
		search,
		"\r\nFilter: ",
		filter,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	Events
//		Control Event Flow.
//		Enable/Disable sending of events to this manager client.
//
func Events(socket *Socket, actionID, eventMask string) (map[string]string, error) {
	if len(actionID) == 0 || len(eventMask) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Events",
		"\r\nActionID: ",
		actionID,
		"\r\nEventMask: ",
		eventMask,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	GetConfig
//		Retrieve configuration.
//		This action will dump the contents of a configuration file by category and contents or optionally by specified category only.
//
func GetConfig(socket *Socket, actionID, filename, category string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 || len(category) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: GetConfig",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\nCategory: ",
		category,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	GetConfigJSON
//		Retrieve configuration (JSON format).
//		This action will dump the contents of a configuration file by category and contents in JSON format.
//		This only makes sense to be used using rawman over the HTTP interface.
//
func GetConfigJSON(socket *Socket, actionID, filename string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: GetConfigJSON",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	JabberSend
//		Sends a message to a Jabber Client.
//
func JabberSend(socket *Socket, actionID, jabber, jid, message string) (map[string]string, error) {
	if len(actionID) == 0 || len(jabber) == 0 ||
		len(jid) == 0 || len(message) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: JabberSend",
		"\r\nActionID: ",
		actionID,
		"\r\nJabber: ",
		jabber,
		"\r\nJID: ",
		jid,
		"\r\nMessage: ",
		message,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ListCommands
//		List available manager commands.
//		Returns the action name and synopsis for every action that is available to the user
//
func ListCommands(socket *Socket, actionID string) (map[string]string, error) {
	if len(actionID) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ListCommands",
		"\r\nActionID: ",
		actionID,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ListCategories
//		List categories in configuration file.
//
func ListCategories(socket *Socket, actionID, filename string) (map[string]string, error) {
	if len(actionID) == 0 || len(filename) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ListCategories",
		"\r\nActionID: ",
		actionID,
		"\r\nFilename: ",
		filename,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ModuleCheck
//		Check if module is loaded.
//		Checks if Asterisk module is loaded. Will return Success/Failure. For success returns, the module revision number is included.
//
func ModuleCheck(socket *Socket, actionID, module string) (map[string]string, error) {
	if len(actionID) == 0 || len(module) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ModuleCheck",
		"\r\nActionID: ",
		actionID,
		"\r\nModule: ",
		module,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ModuleLoad
//		Module management.
//		Loads, unloads or reloads an Asterisk module in a running system.
//
func ModuleLoad(socket *Socket, actionID, module, loadType string) (map[string]string, error) {
	if len(actionID) == 0 || len(module) == 0 || len(loadType) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ModuleLoad",
		"\r\nActionID: ",
		actionID,
		"\r\nModule: ",
		module,
		"\r\nLoadType: ",
		loadType,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	Reload
//		Send a reload event.
//
func Reload(socket *Socket, actionID, module string) (map[string]string, error) {
	if len(actionID) == 0 || len(module) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: Reload",
		"\r\nActionID: ",
		actionID,
		"\r\nModule: ",
		module,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	ShowDialPlan
//		Show dialplan contexts and extensions
//		Be aware that showing the full dialplan may take a lot of capacity
//
func ShowDialPlan(socket *Socket, actionID, extension, context string) (map[string]string, error) {
	if len(actionID) == 0 || len(extension) == 0 || len(context) == 0 {
		return nil, errors.New("Invalid parameters")
	}
	command := []string{
		"Action: ShowDialPlan",
		"\r\nActionID: ",
		actionID,
		"\r\nExtension: ",
		extension,
		"\r\nContext: ",
		context,
		"\r\n\r\n", // end of command
	}
	return getMessage(socket, command, actionID)
}

//	GetEvents
//		Get events from current socket connection
//		It is mandatory set 'events' of ami.Login with "system,call,all,user", to received
//		events
//
func GetEvents(socket *Socket) (map[string]string, error) {
	message, err := decode(socket)
	if err != nil {
		return nil, err
	}
	return message, nil
}
<|endoftext|>


<commit_before>package editor

import (
	"io"
	"sync"
	"unicode/utf8"
)

type RuneAddReader struct {
	ird io.RuneReader

	mu sync.Mutex
	ch chan rune
}

func NewReader(rd io.RuneReader) *RuneAddReader {
	return &RuneAddReader{ird: rd, ch: make(chan rune, 2)}
}

func (rd *RuneAddReader) ReadRune() (r rune, size int, err error) {
	rd.mu.Lock()
	defer rd.mu.Unlock()
	select {
	case r = <-rd.ch:
		size = utf8.RuneLen(r)
		return r, size, err
	default:
		return rd.ird.ReadRune()
	}
}

func (rd *RuneAddReader) Add(s []rune) {
	rd.mu.Lock()
	go func() {
		for _, r := range s {
			rd.ch <- r
		}
		rd.mu.Unlock()
	}()
}
<commit_msg>Don't limit capacity of RuneAddReader to 2<commit_after>package editor

import (
	"io"
	"sync"
	"unicode/utf8"
)

type RuneAddReader struct {
	ird io.RuneReader
	ch  chan rune
	wg  *sync.WaitGroup
}

func NewReader(rd io.RuneReader) *RuneAddReader {
	return &RuneAddReader{ird: rd, ch: make(chan rune), wg: &sync.WaitGroup{}}
}

func (rd *RuneAddReader) ReadRune() (r rune, size int, err error) {
	done := make(chan struct{})
	go func() {
		rd.wg.Wait()
		select {
		case <-done:
			return
		default:
		}
		r, size, err = rd.ird.ReadRune()
		close(done)
	}()
	select {
	case r = <-rd.ch:
		close(done)
		size = utf8.RuneLen(r)
	case <-done:
	}
	return r, size, err
}

func (rd *RuneAddReader) Add(s []rune) {
	rd.wg.Add(len(s))
	go func() {
		for _, r := range s {
			rd.ch <- r
			rd.wg.Done()
		}
	}()
}
<|endoftext|>


<commit_before>package ezk

import (
	cv "github.com/glycerine/goconvey/convey"
	zook "github.com/samuel/go-zookeeper/zk"
	"testing"
	"time"
)

func Test001ClientRetryGetsDefault(t *testing.T) {
	cv.Convey("Given that we don't configure a Retry function, we should get the DefaultRetry function in our ClientConfig", t, func() {

		cli := NewClient(ClientConfig{})
		cv.So(cli.Cfg.Retry, cv.ShouldEqual, DefaultRetry)

	})
}

// In the example code, the goal is to set
// the value newURL into the node /chroot/service-name/config/server-url-list
func Example() {
	newURL := "http://my-new-url.org:343/hello/enhanced-zookeeper-client"

	base := "/chroot"
	nsTest := "/service-name"
	subdir := "/config"
	path := nsTest + subdir + "/server-url-list"
	zkCfg := ClientConfig{
		Servers:        []string{"127.0.0.1:2181"},
		Acl:            zook.WorldACL(zook.PermAll),
		Chroot:         base,
		SessionTimeout: 10 * time.Second,
	}
	zk := NewClient(zkCfg)
	err := zk.Connect()
	if err != nil {
		panic(err)
	}

	defer zk.Close()

	zk.CreateNode(nsTest)
	zk.CreateNode(nsTest + subdir)
	zk.CreateNode(path)

	err = zk.DeleteNode(path) // delete any old value
	if err != nil {
		panic(err)
	}

	err = zk.CreateNode(path)
	if err != nil {
		panic(err)
	}

	_, err = zk.Set(path, []byte(newURL), -1)
	if err != nil {
		panic(err)
	}
}
<commit_msg>atg. include Client example<commit_after>package ezk

import (
	cv "github.com/glycerine/goconvey/convey"
	zook "github.com/samuel/go-zookeeper/zk"
	"testing"
	"time"
)

func Test001ClientRetryGetsDefault(t *testing.T) {
	cv.Convey("Given that we don't configure a Retry function, we should get the DefaultRetry function in our ClientConfig", t, func() {

		cli := NewClient(ClientConfig{})
		cv.So(cli.Cfg.Retry, cv.ShouldEqual, DefaultRetry)

	})
}

// In the example code, the goal is to set
// the value newURL into the node /chroot/service-name/config/server-url-list
func ExampleClient() {
	newURL := "http://my-new-url.org:343/hello/enhanced-zookeeper-client"

	base := "/chroot"
	nsTest := "/service-name"
	subdir := "/config"
	path := nsTest + subdir + "/server-url-list"
	zkCfg := ClientConfig{
		Servers:        []string{"127.0.0.1:2181"},
		Acl:            zook.WorldACL(zook.PermAll),
		Chroot:         base,
		SessionTimeout: 10 * time.Second,
	}
	zk := NewClient(zkCfg)
	err := zk.Connect()
	if err != nil {
		panic(err)
	}

	defer zk.Close()

	zk.CreateNode(nsTest)
	zk.CreateNode(nsTest + subdir)
	zk.CreateNode(path)

	err = zk.DeleteNode(path) // delete any old value
	if err != nil {
		panic(err)
	}

	err = zk.CreateNode(path)
	if err != nil {
		panic(err)
	}

	_, err = zk.Set(path, []byte(newURL), -1)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package curator_test

import (
	. "github.com/talbright/go-curator"
	"github.com/talbright/go-zookeeper/zk"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"path"
	"time"
)

var _ = Describe("Client", func() {
	Context("Connect", func() {
		var client *Client
		AfterEach(func() {
			client.Close()
		})
		It("should connect to server", func() {
			client = NewClient()
			settings := &Settings{
				ZkServers:        zkHosts,
				ZkSessionTimeout: zkSessionTimeout,
			}
			chn, err := client.Connect(settings, zk.WithLogger(&NullLogger{}))
			Expect(err).Should(BeNil())
			events := zkCollectEvents(3, chn)
			Expect(events[0]).Should(MatchEvent(zk.EventSession, zk.StateConnecting))
			Expect(events[1]).Should(MatchEvent(zk.EventSession, zk.StateConnected))
			Expect(events[2]).Should(MatchEvent(zk.EventSession, zk.StateHasSession))
		})
		It("should connect to server and wait for session", func() {
			client = NewClient()
			settings := &Settings{
				ZkServers:               []string{"128.0.0.1:2222"},
				ZkSessionTimeout:        zkSessionTimeout,
				ZkWaitForSessionTimeout: 5 * time.Second,
				ZkWaitForSession:        true,
			}
			_, err := client.Connect(settings, zk.WithLogger(&NullLogger{}))
			Expect(err).Should(Equal(ErrConnectionTimedOut))
		})
	})
	Context("CreatePath", func() {
		var client *Client
		BeforeEach(func() {
			client = zkConnect()
		})
		AfterEach(func() {
			client.Close()
		})
		It("should create the new path /a/b/c", func() {
			basePath := fmt.Sprintf("/test/%s", createUUID())
			patha := path.Join(basePath, "a")
			pathb := path.Join(basePath, "a/b")
			pathc := path.Join(basePath, "a/b/c")
			err := client.CreatePath(pathc, []byte("abc"), zk.WorldACLPermAll)
			Expect(err).Should(BeNil())
			Expect(client).Should(HaveZnode(pathc))
			Expect(client).Should(HaveZnodeData(patha, []byte("abc")))
			acl, _, _ := client.GetACL(patha)
			Expect(acl).Should(Equal(zk.WorldACLPermAll))
			Expect(client).Should(HaveZnodeData(pathb, []byte("abc")))
			acl, _, _ = client.GetACL(pathb)
			Expect(acl).Should(Equal(zk.WorldACLPermAll))
			Expect(client).Should(HaveZnodeData(pathc, []byte("abc")))
			acl, _, _ = client.GetACL(pathc)
			Expect(acl).Should(Equal(zk.WorldACLPermAll))
		})
		It("should create the path /a/b/c/d/e on top of /a/b/c", func() {
			basePath := fmt.Sprintf("/test/%s", createUUID())
			pathc := path.Join(basePath, "a/b/c")
			pathe := path.Join(basePath, "a/b/c/d/e")
			err := client.CreatePath(pathc, []byte("abc"), zk.WorldACLPermAll)
			Expect(err).Should(BeNil())
			err = client.CreatePath(pathe, []byte("de"), zk.WorldACLPermAll)
			Expect(err).Should(BeNil())
			Expect(client).Should(HaveZnode(pathe))
			Expect(client).Should(HaveZnodeData(pathe, []byte("de")))
			Expect(client).Should(HaveZnodeData(pathc, []byte("abc")))
		})
		It("should be an error if the path does not start with a /", func() {
			err := client.CreatePath("a/b/c", zk.NoData, zk.WorldACLPermAll)
			Expect(err).Should(Equal(ErrInvalidPath))
		})
	})
})
<commit_msg>Fix test double closing connection<commit_after>package curator_test

import (
	. "github.com/talbright/go-curator"
	"github.com/talbright/go-zookeeper/zk"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"path"
	"time"
)

var _ = Describe("Client", func() {
	Context("Connect", func() {
		var client *Client
		It("should connect to server", func() {
			client = NewClient()
			settings := &Settings{
				ZkServers:        zkHosts,
				ZkSessionTimeout: zkSessionTimeout,
			}
			chn, err := client.Connect(settings, zk.WithLogger(&NullLogger{}))
			Expect(err).Should(BeNil())
			events := zkCollectEvents(3, chn)
			Expect(events[0]).Should(MatchEvent(zk.EventSession, zk.StateConnecting))
			Expect(events[1]).Should(MatchEvent(zk.EventSession, zk.StateConnected))
			Expect(events[2]).Should(MatchEvent(zk.EventSession, zk.StateHasSession))
			client.Close()
		})
		It("should connect to server and wait for session", func() {
			client = NewClient()
			settings := &Settings{
				ZkServers:               []string{"128.0.0.1:2222"},
				ZkSessionTimeout:        zkSessionTimeout,
				ZkWaitForSessionTimeout: 5 * time.Second,
				ZkWaitForSession:        true,
			}
			_, err := client.Connect(settings, zk.WithLogger(&NullLogger{}))
			Expect(err).Should(Equal(ErrConnectionTimedOut))
		})
	})
	Context("CreatePath", func() {
		var client *Client
		BeforeEach(func() {
			client = zkConnect()
		})
		AfterEach(func() {
			client.Close()
		})
		It("should create the new path /a/b/c", func() {
			basePath := fmt.Sprintf("/test/%s", createUUID())
			patha := path.Join(basePath, "a")
			pathb := path.Join(basePath, "a/b")
			pathc := path.Join(basePath, "a/b/c")
			err := client.CreatePath(pathc, []byte("abc"), zk.WorldACLPermAll)
			Expect(err).Should(BeNil())
			Expect(client).Should(HaveZnode(pathc))
			Expect(client).Should(HaveZnodeData(patha, []byte("abc")))
			acl, _, _ := client.GetACL(patha)
			Expect(acl).Should(Equal(zk.WorldACLPermAll))
			Expect(client).Should(HaveZnodeData(pathb, []byte("abc")))
			acl, _, _ = client.GetACL(pathb)
			Expect(acl).Should(Equal(zk.WorldACLPermAll))
			Expect(client).Should(HaveZnodeData(pathc, []byte("abc")))
			acl, _, _ = client.GetACL(pathc)
			Expect(acl).Should(Equal(zk.WorldACLPermAll))
		})
		It("should create the path /a/b/c/d/e on top of /a/b/c", func() {
			basePath := fmt.Sprintf("/test/%s", createUUID())
			pathc := path.Join(basePath, "a/b/c")
			pathe := path.Join(basePath, "a/b/c/d/e")
			err := client.CreatePath(pathc, []byte("abc"), zk.WorldACLPermAll)
			Expect(err).Should(BeNil())
			err = client.CreatePath(pathe, []byte("de"), zk.WorldACLPermAll)
			Expect(err).Should(BeNil())
			Expect(client).Should(HaveZnode(pathe))
			Expect(client).Should(HaveZnodeData(pathe, []byte("de")))
			Expect(client).Should(HaveZnodeData(pathc, []byte("abc")))
		})
		It("should be an error if the path does not start with a /", func() {
			err := client.CreatePath("a/b/c", zk.NoData, zk.WorldACLPermAll)
			Expect(err).Should(Equal(ErrInvalidPath))
		})
	})
})
<|endoftext|>


<commit_before>package radius

import (
	"context"
	"net"
	"testing"
	"time"
)

func TestClient_Exchange_expired(t *testing.T) {
	handler := HandlerFunc(func(w ResponseWriter, r *Request) {
		// ignore
	})
	server := newTestServer(handler, StaticSecretSource([]byte(`12345`)))
	defer server.Close()

	ctx, cancel := context.WithTimeout(context.Background(), -time.Hour)
	defer cancel()

	req := New(CodeAccessRequest, []byte(`secret`))

	client := Client{}
	resp, err := client.Exchange(ctx, req, server.Addr)
	if resp != nil {
		t.Fatalf("got non-nil response (%v); expected nil", resp)
	}
	if err == nil {
		t.Fatal("got nil error; expected one")
	}
	if netErr, ok := err.(net.Error); !ok {
		t.Fatal("err is not a net.Error")
	} else if !netErr.Timeout() {
		t.Fatal("got netErr.Timeout() = false; expected true")
	}
}
<commit_msg>add test for Client.Retry<commit_after>package radius

import (
	"context"
	"net"
	"sync/atomic"
	"testing"
	"time"
)

func TestClient_Exchange_expired(t *testing.T) {
	handler := HandlerFunc(func(w ResponseWriter, r *Request) {
		// ignore
	})
	server := newTestServer(handler, StaticSecretSource([]byte(`12345`)))
	defer server.Close()

	ctx, cancel := context.WithTimeout(context.Background(), -time.Hour)
	defer cancel()

	req := New(CodeAccessRequest, []byte(`secret`))

	client := Client{}
	resp, err := client.Exchange(ctx, req, server.Addr)
	if resp != nil {
		t.Fatalf("got non-nil response (%v); expected nil", resp)
	}
	if err == nil {
		t.Fatal("got nil error; expected one")
	}
	if netErr, ok := err.(net.Error); !ok {
		t.Fatal("err is not a net.Error")
	} else if !netErr.Timeout() {
		t.Fatal("got netErr.Timeout() = false; expected true")
	}
}

func TestClient_Exchange_retry(t *testing.T) {
	secret := []byte(`12345`)

	var attempts int32

	handler := HandlerFunc(func(w ResponseWriter, r *Request) {
		if atomic.AddInt32(&attempts, 1) == 4 {
			w.Write(r.Response(CodeAccessAccept))
		}
	})
	server := newTestServer(handler, StaticSecretSource(secret))
	defer server.Close()

	req := New(CodeAccessRequest, secret)

	client := Client{
		Retry: time.Millisecond * 5,
	}
	resp, err := client.Exchange(context.Background(), req, server.Addr)
	if err != nil {
		t.Fatalf("got err %s; expected nil", err)
	}
	if resp.Code != CodeAccessAccept {
		t.Fatalf("got code %s; expected %s", resp.Code, CodeAccessAccept)
	}
	if attempts := atomic.LoadInt32(&attempts); attempts != 4 {
		t.Fatalf("response received in %d attemps; expecting 4", attempts)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017-2019 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 cloudprober provides a prober for running a set of probes.

Cloudprober takes in a config proto which dictates what probes should be created
with what configuration, and manages the asynchronous fan-in/fan-out of the
metrics data from these probes.
*/
package cloudprober

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"

	"github.com/golang/protobuf/proto"
	"github.com/google/cloudprober/config"
	configpb "github.com/google/cloudprober/config/proto"
	"github.com/google/cloudprober/config/runconfig"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/prober"
	"github.com/google/cloudprober/probes"
	"github.com/google/cloudprober/servers"
	"github.com/google/cloudprober/surfacers"
	"github.com/google/cloudprober/sysvars"
	"google.golang.org/grpc"
)

const (
	sysvarsModuleName = "sysvars"
)

// Constants defining the default server host and port.
const (
	DefaultServerHost = ""
	DefaultServerPort = 9313
	ServerHostEnvVar  = "CLOUDPROBER_HOST"
	ServerPortEnvVar  = "CLOUDPROBER_PORT"
)

// Global prober.Prober instance protected by a mutex.
var cloudProber struct {
	prober          *prober.Prober
	defaultServerLn net.Listener
	defaultGRPCLn   net.Listener
	textConfig      string
	cancelInitCtx   context.CancelFunc
	sync.Mutex
}

func getServerHost(c *configpb.ProberConfig) string {
	serverHost := c.GetHost()
	if serverHost == "" {
		serverHost = DefaultServerHost
		// If ServerHostEnvVar is defined, it will override the default
		// server host.
		if host := os.Getenv(ServerHostEnvVar); host != "" {
			serverHost = host
		}
	}
	return serverHost
}

func initDefaultServer(c *configpb.ProberConfig) (net.Listener, error) {
	serverHost := getServerHost(c)
	serverPort := int(c.GetPort())
	if serverPort == 0 {
		serverPort = DefaultServerPort

		// If ServerPortEnvVar is defined, it will override the default
		// server port.
		if portStr := os.Getenv(ServerPortEnvVar); portStr != "" {
			port, err := strconv.ParseInt(portStr, 10, 32)
			if err != nil {
				return nil, fmt.Errorf("failed to parse default port from the env var: %s=%s", ServerPortEnvVar, portStr)
			}
			serverPort = int(port)
		}
	}

	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", serverHost, serverPort))
	if err != nil {
		return nil, fmt.Errorf("error while creating listener for default HTTP server: %v", err)
	}

	return ln, nil
}

// InitFromConfig initializes Cloudprober using the provided config.
func InitFromConfig(configFile string) error {
	// Return immediately if prober is already initialized.
	cloudProber.Lock()
	defer cloudProber.Unlock()

	if cloudProber.prober != nil {
		return nil
	}

	// Initialize sysvars module
	l, err := logger.NewCloudproberLog(sysvarsModuleName)
	if err != nil {
		return err
	}
	sysvars.Init(l, nil)

	configStr, err := config.ParseTemplate(configFile, sysvars.Vars())
	if err != nil {
		return err
	}

	cfg := &configpb.ProberConfig{}
	if err := proto.UnmarshalText(configStr, cfg); err != nil {
		return err
	}

	globalLogger, err := logger.NewCloudproberLog("global")
	if err != nil {
		return fmt.Errorf("error in initializing global logger: %v", err)
	}

	// Start default HTTP server. It's used for profile handlers and
	// prometheus exporter.
	ln, err := initDefaultServer(cfg)
	if err != nil {
		return err
	}

	var grpcLn net.Listener
	if cfg.GetGrpcPort() != 0 {
		serverHost := getServerHost(cfg)

		grpcLn, err = net.Listen("tcp", fmt.Sprintf("%s:%d", serverHost, cfg.GetGrpcPort()))
		if err != nil {
			return fmt.Errorf("error while creating listener for default gRPC server: %v", err)
		}

		// Create the default gRPC server now, so that other modules can register
		// their services with it in the prober.Init() phase.
		runconfig.SetDefaultGRPCServer(grpc.NewServer())
	}

	pr := &prober.Prober{}

	// initCtx is used to clean up in case of partial initialization failures. For
	// example, user-configured servers open listeners during initialization and
	// if initialization fails at a later stage, say in probers or surfacers,
	// pr.Init returns an error and we cancel the initCtx, which makes servers
	// close their listeners.
	// TODO(manugarg): Plumb init context from cmd/cloudprober.
	initCtx, cancelFunc := context.WithCancel(context.TODO())
	if err := pr.Init(initCtx, cfg, globalLogger); err != nil {
		cancelFunc()
		ln.Close()
		return err
	}

	cloudProber.prober = pr
	cloudProber.textConfig = configStr
	cloudProber.defaultServerLn = ln
	cloudProber.defaultGRPCLn = grpcLn
	cloudProber.cancelInitCtx = cancelFunc
	return nil
}

// Start starts a previously initialized Cloudprober.
func Start(ctx context.Context) {
	cloudProber.Lock()
	defer cloudProber.Unlock()

	// Default servers
	srv := &http.Server{}
	grpcSrv := runconfig.DefaultGRPCServer()

	// Set up a goroutine to cleanup if context ends.
	go func() {
		<-ctx.Done()
		srv.Close() // This will close the listener as well.
		if grpcSrv != nil {
			grpcSrv.Stop()
		}
		cloudProber.cancelInitCtx()
	}()

	go srv.Serve(cloudProber.defaultServerLn)
	if grpcSrv != nil {
		go grpcSrv.Serve(cloudProber.defaultGRPCLn)
	}

	if cloudProber.prober == nil {
		panic("Prober is not initialized. Did you call cloudprober.InitFromConfig first?")
	}

	cloudProber.prober.Start(ctx)
}

// GetConfig returns the prober config.
func GetConfig() string {
	cloudProber.Lock()
	defer cloudProber.Unlock()
	return cloudProber.textConfig
}

// GetInfo returns information on all the probes, servers and surfacers.
func GetInfo() (map[string]*probes.ProbeInfo, []*surfacers.SurfacerInfo, []*servers.ServerInfo) {
	cloudProber.Lock()
	defer cloudProber.Unlock()
	return cloudProber.prober.Probes, cloudProber.prober.Surfacers, cloudProber.prober.Servers
}
<commit_msg>Stop cloudprober from starting gRPC server if a grpc listener is not created.<commit_after>// Copyright 2017-2019 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 cloudprober provides a prober for running a set of probes.

Cloudprober takes in a config proto which dictates what probes should be created
with what configuration, and manages the asynchronous fan-in/fan-out of the
metrics data from these probes.
*/
package cloudprober

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"

	"github.com/golang/protobuf/proto"
	"github.com/google/cloudprober/config"
	configpb "github.com/google/cloudprober/config/proto"
	"github.com/google/cloudprober/config/runconfig"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/prober"
	"github.com/google/cloudprober/probes"
	"github.com/google/cloudprober/servers"
	"github.com/google/cloudprober/surfacers"
	"github.com/google/cloudprober/sysvars"
	"google.golang.org/grpc"
)

const (
	sysvarsModuleName = "sysvars"
)

// Constants defining the default server host and port.
const (
	DefaultServerHost = ""
	DefaultServerPort = 9313
	ServerHostEnvVar  = "CLOUDPROBER_HOST"
	ServerPortEnvVar  = "CLOUDPROBER_PORT"
)

// Global prober.Prober instance protected by a mutex.
var cloudProber struct {
	prober          *prober.Prober
	defaultServerLn net.Listener
	defaultGRPCLn   net.Listener
	textConfig      string
	cancelInitCtx   context.CancelFunc
	sync.Mutex
}

func getServerHost(c *configpb.ProberConfig) string {
	serverHost := c.GetHost()
	if serverHost == "" {
		serverHost = DefaultServerHost
		// If ServerHostEnvVar is defined, it will override the default
		// server host.
		if host := os.Getenv(ServerHostEnvVar); host != "" {
			serverHost = host
		}
	}
	return serverHost
}

func initDefaultServer(c *configpb.ProberConfig) (net.Listener, error) {
	serverHost := getServerHost(c)
	serverPort := int(c.GetPort())
	if serverPort == 0 {
		serverPort = DefaultServerPort

		// If ServerPortEnvVar is defined, it will override the default
		// server port.
		if portStr := os.Getenv(ServerPortEnvVar); portStr != "" {
			port, err := strconv.ParseInt(portStr, 10, 32)
			if err != nil {
				return nil, fmt.Errorf("failed to parse default port from the env var: %s=%s", ServerPortEnvVar, portStr)
			}
			serverPort = int(port)
		}
	}

	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", serverHost, serverPort))
	if err != nil {
		return nil, fmt.Errorf("error while creating listener for default HTTP server: %v", err)
	}

	return ln, nil
}

// InitFromConfig initializes Cloudprober using the provided config.
func InitFromConfig(configFile string) error {
	// Return immediately if prober is already initialized.
	cloudProber.Lock()
	defer cloudProber.Unlock()

	if cloudProber.prober != nil {
		return nil
	}

	// Initialize sysvars module
	l, err := logger.NewCloudproberLog(sysvarsModuleName)
	if err != nil {
		return err
	}
	sysvars.Init(l, nil)

	configStr, err := config.ParseTemplate(configFile, sysvars.Vars())
	if err != nil {
		return err
	}

	cfg := &configpb.ProberConfig{}
	if err := proto.UnmarshalText(configStr, cfg); err != nil {
		return err
	}

	globalLogger, err := logger.NewCloudproberLog("global")
	if err != nil {
		return fmt.Errorf("error in initializing global logger: %v", err)
	}

	// Start default HTTP server. It's used for profile handlers and
	// prometheus exporter.
	ln, err := initDefaultServer(cfg)
	if err != nil {
		return err
	}

	var grpcLn net.Listener
	if cfg.GetGrpcPort() != 0 {
		serverHost := getServerHost(cfg)

		grpcLn, err = net.Listen("tcp", fmt.Sprintf("%s:%d", serverHost, cfg.GetGrpcPort()))
		if err != nil {
			return fmt.Errorf("error while creating listener for default gRPC server: %v", err)
		}

		// Create the default gRPC server now, so that other modules can register
		// their services with it in the prober.Init() phase.
		runconfig.SetDefaultGRPCServer(grpc.NewServer())
	}

	pr := &prober.Prober{}

	// initCtx is used to clean up in case of partial initialization failures. For
	// example, user-configured servers open listeners during initialization and
	// if initialization fails at a later stage, say in probers or surfacers,
	// pr.Init returns an error and we cancel the initCtx, which makes servers
	// close their listeners.
	// TODO(manugarg): Plumb init context from cmd/cloudprober.
	initCtx, cancelFunc := context.WithCancel(context.TODO())
	if err := pr.Init(initCtx, cfg, globalLogger); err != nil {
		cancelFunc()
		ln.Close()
		return err
	}

	cloudProber.prober = pr
	cloudProber.textConfig = configStr
	cloudProber.defaultServerLn = ln
	cloudProber.defaultGRPCLn = grpcLn
	cloudProber.cancelInitCtx = cancelFunc
	return nil
}

// Start starts a previously initialized Cloudprober.
func Start(ctx context.Context) {
	cloudProber.Lock()
	defer cloudProber.Unlock()

	// Default servers
	srv := &http.Server{}
	grpcSrv := runconfig.DefaultGRPCServer()

	// Set up a goroutine to cleanup if context ends.
	go func() {
		<-ctx.Done()
		srv.Close() // This will close the listener as well.
		if grpcSrv != nil {
			grpcSrv.Stop()
		}
		cloudProber.cancelInitCtx()
	}()

	go srv.Serve(cloudProber.defaultServerLn)
	if grpcSrv != nil && cloudProber.defaultGRPCLn != nil {
		go grpcSrv.Serve(cloudProber.defaultGRPCLn)
	}

	if cloudProber.prober == nil {
		panic("Prober is not initialized. Did you call cloudprober.InitFromConfig first?")
	}

	cloudProber.prober.Start(ctx)
}

// GetConfig returns the prober config.
func GetConfig() string {
	cloudProber.Lock()
	defer cloudProber.Unlock()
	return cloudProber.textConfig
}

// GetInfo returns information on all the probes, servers and surfacers.
func GetInfo() (map[string]*probes.ProbeInfo, []*surfacers.SurfacerInfo, []*servers.ServerInfo) {
	cloudProber.Lock()
	defer cloudProber.Unlock()
	return cloudProber.prober.Probes, cloudProber.prober.Surfacers, cloudProber.prober.Servers
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/avarabyeu/releaser/util"
	"github.com/juju/errgo/errors"
	"github.com/spf13/cobra"
	"gopkg.in/cheggaaa/pb.v1"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"sync"
	"time"
)

const (
	refreshRate  = time.Millisecond * 10
	btVersionURL = "https://api.bintray.com/packages/{{.org}}/{{.repo}}/{{.package}}/versions"
)

var btVersionTemplate *template.Template

func init() {
	btVersionTemplate, _ = template.New("btVersionTMPL").Parse(btVersionURL)
}

var bintrayCommand = &cobra.Command{
	Use:   "bintray",
	Short: "Uploads to bintray",
	Long:  `Uploads artifacts to bintray`,
	Run: func(cmd *cobra.Command, args []string) {
		err := uploadToBintray(cmd)
		if nil != err {
			log.Fatalf("Cannot upload to bintray. %s", err)
		}
	},
}

func uploadToBintray(cmd *cobra.Command) error {
	wd, err := os.Getwd()
	if nil != err {
		log.Printf("Cannot get work dir! %s", err.Error())
		return err
	}

	err = checkData(config.Bintray)
	if nil != err {
		log.Printf("BintrayConf data isn't provided: %s", err.Error())
		return err
	}

	artifactsFolder := path.Join(wd, cmd.Flag("artifactsFolder").Value.String())
	files, _ := ioutil.ReadDir(artifactsFolder)

	if len(files) > 0 {
		er := createVersion(GetSemver(cmd).Version.String(), config.Bintray)
		if nil != er {
			return er
		}

	}

	// update bars
	pool, err := pb.StartPool()
	if err != nil {
		return err
	}

	wg := new(sync.WaitGroup)
	for _, f := range files {
		wg.Add(1)
		log.Print(f.Name())
		go uploadFile(GetSemver(cmd).Current(), path.Join(artifactsFolder, f.Name()), config.Bintray, wg, pool)
	}

	wg.Wait()
	pool.Stop()
	return nil
}

func createVersion(version string, bd *BintrayConf) error {
	details := map[string]string{"name": version, "desc": "New version"}
	rqBody, err := json.Marshal(details)
	if nil != err {
		return err
	}

	url := new(bytes.Buffer)
	btVersionTemplate.Execute(url, map[string]string{"org": bd.Org, "repo": bd.Repo, "package": bd.Pack})

	log.Println(url.String())
	rq, err := http.NewRequest("POST", url.String(), bytes.NewBuffer(rqBody))
	if err != nil {
		return err
	}
	rq.Header.Set("Authorization", bd.getAuth())
	rq.Header.Set("Content-Type", "application/json")

	rs, err := http.DefaultClient.Do(rq)
	if err != nil {
		return err
	}
	defer rs.Body.Close()

	if rs.StatusCode/100 != 2 {
		body, _ := ioutil.ReadAll(rs.Body)
		fmt.Println("Error:", string(body))
		return errors.Newf("Cannot create new version. Status code: %d", rs.StatusCode)
	}

	return nil
}

func uploadFile(version string, filePath string, bt *BintrayConf, wg *sync.WaitGroup, pool *pb.Pool) {
	defer wg.Done()
	var err error
	var f *os.File
	var fi os.FileInfo

	var bar *pb.ProgressBar

	if f, err = os.Open(filePath); err != nil {
		log.Fatal(err)
	}
	if fi, err = f.Stat(); err != nil {
		log.Fatal(err)
	}
	bar = pb.New64(fi.Size()).Prefix(util.SubstrAfterLast(f.Name(), "/")).SetUnits(pb.U_BYTES).SetRefreshRate(refreshRate)
	pool.Add(bar)

	r, w := io.Pipe()
	go func() {
		defer w.Close()
		defer f.Close()

		rqw := io.MultiWriter(w, bar)
		if _, err = io.Copy(rqw, f); err != nil {
			log.Fatal(err)
		}

	}()

	rq, err := http.NewRequest("PUT", fmt.Sprintf("https://api.bintray.com/content/%s/%s/%s/%s", bt.Org, bt.Repo, version, fi.Name()), r)
	rq.Header.Set("Content-Type", "application/octet-stream")
	rq.Header.Set("X-Bintray-Package", bt.Pack)
	rq.Header.Set("X-Bintray-Version", version)
	rq.Header.Set("X-Bintray-Override", "1")
	rq.Header.Set("X-Bintray-Publish", "1")
	rq.Header.Set("Authorization", bt.getAuth())

	resp, err := http.DefaultClient.Do(rq)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()
	rs, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	} else {
		log.Println(string(rs))
		log.Println(resp.StatusCode)
	}
}

func checkData(btData *BintrayConf) error {
	var errs string

	if "" == btData.Org {
		errs += "BintrayConf organization not set\n"
	}

	if "" == btData.User {
		errs += "BintrayConf user not set\n"
	}

	if "" == btData.Token {
		errs += "BintrayConf token not set\n"
	}

	if "" == btData.Pack {
		errs += "BintrayConf package not set\n"
	}

	if "" == btData.Repo {
		errs += "BintrayConf repo not set\n"
	}

	if "" != errs {
		log.Fatalf("Validation failed: \n%s", errs)
		return errors.New(errs)
	}

	return nil

}
<commit_msg>temparaly disable pool progress bar<commit_after>package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/avarabyeu/releaser/util"
	"github.com/juju/errgo/errors"
	"github.com/spf13/cobra"
	"gopkg.in/cheggaaa/pb.v1"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"time"
)

const (
	refreshRate  = time.Millisecond * 10
	btVersionURL = "https://api.bintray.com/packages/{{.org}}/{{.repo}}/{{.package}}/versions"
)

var btVersionTemplate *template.Template

func init() {
	btVersionTemplate, _ = template.New("btVersionTMPL").Parse(btVersionURL)
}

var bintrayCommand = &cobra.Command{
	Use:   "bintray",
	Short: "Uploads to bintray",
	Long:  `Uploads artifacts to bintray`,
	Run: func(cmd *cobra.Command, args []string) {
		err := uploadToBintray(cmd)
		if nil != err {
			log.Fatalf("Cannot upload to bintray. %s", err)
		}
	},
}

func uploadToBintray(cmd *cobra.Command) error {
	wd, err := os.Getwd()
	if nil != err {
		log.Printf("Cannot get work dir! %s", err.Error())
		return err
	}

	err = checkData(config.Bintray)
	if nil != err {
		log.Printf("BintrayConf data isn't provided: %s", err.Error())
		return err
	}

	artifactsFolder := path.Join(wd, cmd.Flag("artifactsFolder").Value.String())
	files, _ := ioutil.ReadDir(artifactsFolder)

	if len(files) > 0 {
		er := createVersion(GetSemver(cmd).Version.String(), config.Bintray)
		if nil != er {
			return er
		}

	}

	// update bars
	//pool, err := pb.StartPool()
	if err != nil {
		return err
	}

	//wg := new(sync.WaitGroup)
	for _, f := range files {
		//wg.Add(1)
		log.Print(f.Name())
		uploadFile(GetSemver(cmd).Current(), path.Join(artifactsFolder, f.Name()), config.Bintray)
	}

	//wg.Wait()
	//pool.Stop()
	return nil
}

func createVersion(version string, bd *BintrayConf) error {
	details := map[string]string{"name": version, "desc": "New version"}
	rqBody, err := json.Marshal(details)
	if nil != err {
		return err
	}

	url := new(bytes.Buffer)
	btVersionTemplate.Execute(url, map[string]string{"org": bd.Org, "repo": bd.Repo, "package": bd.Pack})

	log.Println(url.String())
	rq, err := http.NewRequest("POST", url.String(), bytes.NewBuffer(rqBody))
	if err != nil {
		return err
	}
	rq.Header.Set("Authorization", bd.getAuth())
	rq.Header.Set("Content-Type", "application/json")

	rs, err := http.DefaultClient.Do(rq)
	if err != nil {
		return err
	}
	defer rs.Body.Close()

	if rs.StatusCode/100 != 2 {
		body, _ := ioutil.ReadAll(rs.Body)
		fmt.Println("Error:", string(body))
		return errors.Newf("Cannot create new version. Status code: %d", rs.StatusCode)
	}

	return nil
}

func uploadFile(version string, filePath string, bt *BintrayConf) {
	//defer wg.Done()
	var err error
	var f *os.File
	var fi os.FileInfo

	var bar *pb.ProgressBar

	if f, err = os.Open(filePath); err != nil {
		log.Fatal(err)
	}
	if fi, err = f.Stat(); err != nil {
		log.Fatal(err)
	}
	bar = pb.New64(fi.Size()).Prefix(util.SubstrAfterLast(f.Name(), "/")).SetUnits(pb.U_BYTES).SetRefreshRate(refreshRate)
	//pool.Add(bar)

	r, w := io.Pipe()
	go func() {
		defer w.Close()
		defer f.Close()

		rqw := io.MultiWriter(w, bar)
		if _, err = io.Copy(rqw, f); err != nil {
			log.Fatal(err)
		}

	}()

	rq, err := http.NewRequest("PUT", fmt.Sprintf("https://api.bintray.com/content/%s/%s/%s/%s", bt.Org, bt.Repo, version, fi.Name()), r)
	rq.Header.Set("Content-Type", "application/octet-stream")
	rq.Header.Set("X-Bintray-Package", bt.Pack)
	rq.Header.Set("X-Bintray-Version", version)
	rq.Header.Set("X-Bintray-Override", "1")
	rq.Header.Set("X-Bintray-Publish", "1")
	rq.Header.Set("Authorization", bt.getAuth())

	resp, err := http.DefaultClient.Do(rq)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()
	rs, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	} else {
		log.Println(string(rs))
		log.Println(resp.StatusCode)
	}
}

func checkData(btData *BintrayConf) error {
	var errs string

	if "" == btData.Org {
		errs += "BintrayConf organization not set\n"
	}

	if "" == btData.User {
		errs += "BintrayConf user not set\n"
	}

	if "" == btData.Token {
		errs += "BintrayConf token not set\n"
	}

	if "" == btData.Pack {
		errs += "BintrayConf package not set\n"
	}

	if "" == btData.Repo {
		errs += "BintrayConf repo not set\n"
	}

	if "" != errs {
		log.Fatalf("Validation failed: \n%s", errs)
		return errors.New(errs)
	}

	return nil

}
<|endoftext|>


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

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/vdobler/ht/ht"
	"github.com/vdobler/ht/internal/bender"
	"github.com/vdobler/ht/suite"
)

var cmdLoad = &Command{
	RunArgs:     runLoad,
	Usage:       "load [options] <loadtest>",
	Description: "run suites in a load test",
	Flag:        flag.NewFlagSet("load", flag.ContinueOnError),
	Help: `
Execute a throughput test with the given suite.
	`,
}

var queryPerSecond float64
var testDuration time.Duration
var rampDuration time.Duration

func init() {
	cmdLoad.Flag.Float64Var(&queryPerSecond, "rate", 20,
		"make `qps` reqest per second")
	cmdLoad.Flag.DurationVar(&testDuration, "duration", 30*time.Second,
		"duration of throughput test")
	cmdLoad.Flag.DurationVar(&rampDuration, "ramp", 5*time.Second,
		"ramp duration to reach desired request rate")
	addOutputFlag(cmdLoad.Flag)
	addVarsFlags(cmdLoad.Flag)
}

func runLoad(cmd *Command, args []string) {
	if len(args) != 1 {
		fmt.Fprintf(os.Stderr, "Usage: %s\n", cmd.Usage)
		os.Exit(9)
	}
	arg := args[0]

	// Process arguments of the form <name>@<archive>.
	var fs suite.FileSystem = nil
	if i := strings.Index(arg, "@"); i != -1 {
		blob, err := ioutil.ReadFile(arg[i+1:])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot load %q: %s\n", arg[i+1:], err)
			os.Exit(9)
		}
		fs, err = suite.NewFileSystem(string(blob))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot load %q: %s\n", arg[i+1:], err)
			os.Exit(9)
		}
		arg = arg[:i]
	}
	raw, err := suite.LoadRawLoadtest(arg, fs)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Cannot load %q: %s\n", arg, err)
		os.Exit(9)
	}

	scenarios := raw.ToScenario(variablesFlag)
	bufferedStdout := bufio.NewWriterSize(os.Stdout, 512)
	defer bufferedStdout.Flush()
	for i, scen := range scenarios {
		fmt.Printf("%d. %d%% %q (max %d threads, verbosity %d)\n",
			i+1, scen.Percentage, scen.RawSuite.Name, scen.MaxThreads,
			scen.Verbosity)
	}

	prepareHT()
	data, failures, lterr := suite.Throughput(scenarios, queryPerSecond, testDuration, rampDuration)

	if outputDir == "" {
		outputDir = time.Now().Format("2006-01-02_15h04m05s")
	}
	err = os.MkdirAll(outputDir, 0766)
	if err != nil {
		log.Panic(err)
	}

	if failures != nil {
		failures.Name = "Failures of throughput test " + arg
		saveLoadtestData(data, failures)
	}
	printStatistics(scenarios, data)
	interpretLTerrors(lterr)
}

func printStatistics(scenarios []suite.Scenario, data []bender.TestData) {
	// Per testfile
	pert := make(map[string][]bender.TestData)
	for _, d := range data {
		parts := strings.Split(d.ID, suite.IDSep)
		nums := strings.Split(parts[0], "/")
		s, _ := strconv.Atoi(nums[0])
		t, _ := strconv.Atoi(nums[3])
		name := scenarios[s-1].RawSuite.RawTests()[t-1].File.Name
		pert[name] = append(pert[name], d)
	}
	for file, fd := range pert {
		st := statsFor(fd)
		h := fmt.Sprintf("Testfile %q:", file)
		printStat(h, st)

	}

	// Per scenario
	for i, s := range scenarios {
		fd := make([]bender.TestData, 0, 200)
		pat := fmt.Sprintf("%d/", i+1)
		for _, d := range data {
			if strings.HasPrefix(d.ID, pat) {
				fd = append(fd, d)
			}
		}

		st := statsFor(fd)
		h := fmt.Sprintf("Scenario %d %q:", i+1, s.Name)
		printStat(h, st)
	}

	// All requests
	st := statsFor(data)
	printStat("All request:", st)

	/*
		// Per (scenario/test)
		for i, s := range scenarios {
			for j, t := range s.RawSuite.RawTests() {
				fd := make([]bender.TestData, 0, 200)
				ppat := fmt.Sprintf("%d/", i+1)
				spat := fmt.Sprintf("/%d%s", j+1, suite.IDSep)
				for _, d := range data {
					// fmt.Printf("ID = %s   %s  %s\n", d.ID, ppat, spat)
					if strings.HasPrefix(d.ID, ppat) &&
						strings.Contains(d.ID, spat) {
						fd = append(fd, d)
					}
				}
				st := statsFor(fd)
				fmt.Printf("Scenario %d %q, Test %d %q:\n",
					i+1, s.Name,
					j+1, t.File.Name,
				)
				printStat(st)
			}
		}
	*/

}

func printStat(headline string, st sdata) {
	fmt.Printf("%s Status:   Total=%d  Pass=%d (%.1f%%), Fail=%d (%.1f%%), Error=%d (%.1f%%), Bogus=%d (%.1f%%)\n",
		headline, st.n,
		st.good, 100*float64(st.good)/float64(st.n),
		st.fail, 100*float64(st.fail)/float64(st.n),
		st.erred, 100*float64(st.erred)/float64(st.n),
		st.bogus, 100*float64(st.bogus)/float64(st.n),
	)
	fmt.Printf("%s Duration: 0%%=%.1fms, 25%%=%.1fms, 50%%=%.1fms, 75%%=%.1fms, 90%%=%.1fms, 95%%=%.1fms, 99%%=%.1fms, 100%%=%.1fms\n",
		headline,
		float64(st.min/1000)/1000,
		float64(st.q25/1000)/1000,
		float64(st.median/1000)/1000,
		float64(st.q75/1000)/1000,
		float64(st.q90/1000)/1000,
		float64(st.q95/1000)/1000,
		float64(st.q99/1000)/1000,
		float64(st.max/1000)/1000,
	)
}

type sdata struct {
	n                       int
	fail, erred, bogus      int
	min, mean, max, median  time.Duration
	q25, q75, q90, q95, q99 time.Duration
	bad, good               int
}

func statsFor(data []bender.TestData) sdata {
	sd := sdata{
		n:    len(data),
		min:  99 * time.Hour,
		mean: time.Duration(0),
		max:  time.Duration(0),
	}
	if sd.n == 0 {
		return sd
	}

	x := make([]time.Duration, len(data))
	sum := time.Duration(0)
	for i, d := range data {
		if d.Status == ht.Fail {
			sd.fail++
		} else if d.Status == ht.Error {
			sd.erred++
		} else if d.Status == ht.Bogus {
			sd.bogus++
		}
		x[i] = d.ReqDuration
		sum += d.ReqDuration
	}
	sd.bad = sd.fail + sd.erred + sd.bogus
	sd.good = sd.n - sd.bad

	sd.mean = sum / time.Duration(len(data))
	sort.Sort(durationSlice(x))
	sd.min, sd.max = x[0], x[len(x)-1]
	sd.median = quantile(x, 0.5)

	sd.q25, sd.median = quantile(x, 0.25), quantile(x, 0.5)
	sd.q75, sd.q90 = quantile(x, 0.75), quantile(x, 0.90)
	sd.q95, sd.q99 = quantile(x, 0.95), quantile(x, 0.99)
	return sd
}

type durationSlice []time.Duration

func (p durationSlice) Len() int           { return len(p) }
func (p durationSlice) Less(i, j int) bool { return p[i] < p[j] }
func (p durationSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func quantile(x []time.Duration, p float64) time.Duration {
	N := float64(len(x))
	if N == 0 {
		return 0
	} else if N == 1 {
		return x[0]
	}

	if p < 2.0/(3.0*(N+1.0/3.0)) {
		return x[0]
	}
	if p >= (N-1.0/3.0)/(N+1.0/3.0) {
		return x[len(x)-1]
	}

	h := (N+1.0/3.0)*p + 1.0/3.0
	fh := math.Floor(h)
	xl := x[int(fh)-1]
	xr := x[int(fh)]

	return time.Duration(float64(xl) + (h-fh)*float64(xr-xl))
}

func interpretLTerrors(lterr error) {
	if lterr == nil {
		fmt.Println("OKAY")
		os.Exit(0)
	}

	fmt.Println("Problems running this throughpout tests:")
	if el, ok := lterr.(ht.ErrorList); ok {
		for _, msg := range el.AsStrings() {
			fmt.Println("    ", msg)
		}
	} else {
		fmt.Println("  ", lterr.Error())
	}

	fmt.Println("PROBLEMS")
	os.Exit(1)
}

func saveLoadtestData(data []bender.TestData, failures *suite.Suite) {
	err := suite.HTMLReport(outputDir, failures)
	if err != nil {
		log.Panic(err)
	}

	file, err := os.Create(outputDir + "/throughput.csv")
	if err != nil {
		log.Panic(err)
	}
	defer file.Close()

	suite.DataToCSV(data, file)
	script := `
library(ggplot2)
d <- read.csv("throughput.csv")
d$Status <- factor(d$Status, levels <- c("NotRun", "Skipped", "Pass", "Fail", "Error", "Bogus"))

shift <- function(x, lag) {
    n <- length(x)
    xnew <- rep(NA, n)
    xnew[(lag+1):n] <- x[1:(n-lag)]
    return(xnew)
}

p <- ggplot(d, aes(x=Elapsed, y=Rate))
p <- p + geom_point(size=3) + geom_smooth()
ggsave("rate.png", plot=p, width=10, height=8, dpi=100)

d$Delta <- c(d$Elapsed[2:length(d$Elapsed)] - d$Elapsed[1:length(d$Elapsed)-1], NA)
p <- ggplot(d, aes(x=Elapsed, y=Delta)) + geom_point(size=2) + xlab("Elapsed [ms]")
ggsave("delta.png", plot=p, width=10, height=8, dpi=100)

myColors <- c("#999999", "#ffff00", "#339900", "#660000", "#ff0000", "#ff3399")
names(myColors) <- levels(d$Status)
colScale <- scale_colour_manual(name = "status",values = myColors)
fillScale <- scale_fill_manual(name = "status",values = myColors)

p <- ggplot(d, aes(x=Elapsed, y=ReqDuration, colour=Test))
p <- p + geom_point(size=3) + xlab("Elapsed [ms]") + ylab("Request Duration [ms]")
ggsave("scatter.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=Elapsed, y=ReqDuration, colour=Status))
p <- p + geom_point(size=3) + xlab("Elapsed [ms]")  + ylab("Request Duration [ms]")
p <- p + colScale
ggsave("status.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=Elapsed, y=ConcTot, colour=Status)) + colScale
p <- p + geom_point(size=3) + xlab("Elapsed [ms]") + ylab("Total Concurrency")
ggsave("conctot.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=Elapsed, y=ConcOwn, colour=Status)) + colScale
p <- p + facet_grid(Test ~ ., scales="free_y")
p <- p + geom_point(size=3) + xlab("Elapsed [ms]")
ggsave("concown.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=ReqDuration, fill=Status))
p <- p + geom_histogram(binwidth=3)
p <- p + facet_grid(Test ~ ., scales="free_y")
p <- p + fillScale
ggsave("hist.png", plot=p, width=10, height=8, dpi=100)

`
	ioutil.WriteFile(outputDir+"/throughput.R", []byte(script), 0666)

}
<commit_msg>cmd/ht: print sparkline histogram in load test<commit_after>// Copyright 2016 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/vdobler/ht/hist"
	"github.com/vdobler/ht/ht"
	"github.com/vdobler/ht/internal/bender"
	"github.com/vdobler/ht/suite"
)

var cmdLoad = &Command{
	RunArgs:     runLoad,
	Usage:       "load [options] <loadtest>",
	Description: "run suites in a load test",
	Flag:        flag.NewFlagSet("load", flag.ContinueOnError),
	Help: `
Execute a throughput test with the given suite.
	`,
}

var queryPerSecond float64
var testDuration time.Duration
var rampDuration time.Duration

func init() {
	cmdLoad.Flag.Float64Var(&queryPerSecond, "rate", 20,
		"make `qps` reqest per second")
	cmdLoad.Flag.DurationVar(&testDuration, "duration", 30*time.Second,
		"duration of throughput test")
	cmdLoad.Flag.DurationVar(&rampDuration, "ramp", 5*time.Second,
		"ramp duration to reach desired request rate")
	addOutputFlag(cmdLoad.Flag)
	addVarsFlags(cmdLoad.Flag)
}

func runLoad(cmd *Command, args []string) {
	if len(args) != 1 {
		fmt.Fprintf(os.Stderr, "Usage: %s\n", cmd.Usage)
		os.Exit(9)
	}
	arg := args[0]

	// Process arguments of the form <name>@<archive>.
	var fs suite.FileSystem = nil
	if i := strings.Index(arg, "@"); i != -1 {
		blob, err := ioutil.ReadFile(arg[i+1:])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot load %q: %s\n", arg[i+1:], err)
			os.Exit(9)
		}
		fs, err = suite.NewFileSystem(string(blob))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot load %q: %s\n", arg[i+1:], err)
			os.Exit(9)
		}
		arg = arg[:i]
	}
	raw, err := suite.LoadRawLoadtest(arg, fs)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Cannot load %q: %s\n", arg, err)
		os.Exit(9)
	}

	scenarios := raw.ToScenario(variablesFlag)
	bufferedStdout := bufio.NewWriterSize(os.Stdout, 512)
	defer bufferedStdout.Flush()
	for i, scen := range scenarios {
		fmt.Printf("%d. %d%% %q (max %d threads, verbosity %d)\n",
			i+1, scen.Percentage, scen.RawSuite.Name, scen.MaxThreads,
			scen.Verbosity)
	}

	prepareHT()
	data, failures, lterr := suite.Throughput(scenarios, queryPerSecond, testDuration, rampDuration)

	if outputDir == "" {
		outputDir = time.Now().Format("2006-01-02_15h04m05s")
	}
	err = os.MkdirAll(outputDir, 0766)
	if err != nil {
		log.Panic(err)
	}

	if failures != nil {
		failures.Name = "Failures of throughput test " + arg
		saveLoadtestData(data, failures)
	}
	printStatistics(scenarios, data)
	interpretLTerrors(lterr)
}

func printStatistics(scenarios []suite.Scenario, data []bender.TestData) {
	histograms := []hist.Histogram{}

	// Per testfile
	pert := make(map[string][]bender.TestData)
	for _, d := range data {
		parts := strings.Split(d.ID, suite.IDSep)
		nums := strings.Split(parts[0], "/")
		s, _ := strconv.Atoi(nums[0])
		t, _ := strconv.Atoi(nums[3])
		name := scenarios[s-1].RawSuite.RawTests()[t-1].File.Name
		pert[name] = append(pert[name], d)
	}
	for file, fd := range pert {
		st := statsFor(fd)
		h := fmt.Sprintf("Testfile %q:", file)
		printStat(h, st)
		histograms = append(histograms, hist.Histogram{Name: file, Data: st.data})
	}

	// Per scenario
	for i, s := range scenarios {
		fd := make([]bender.TestData, 0, 200)
		pat := fmt.Sprintf("%d/", i+1)
		for _, d := range data {
			if strings.HasPrefix(d.ID, pat) {
				fd = append(fd, d)
			}
		}

		st := statsFor(fd)
		h := fmt.Sprintf("Scenario %d %q:", i+1, s.Name)
		printStat(h, st)
	}

	// All requests
	st := statsFor(data)
	printStat("All request:", st)
	histograms = append(histograms, hist.Histogram{Name: "All requests:", Data: st.data})
	hist.PrintLogHistograms(os.Stdout, histograms)

	/*
		// Per (scenario/test)
		for i, s := range scenarios {
			for j, t := range s.RawSuite.RawTests() {
				fd := make([]bender.TestData, 0, 200)
				ppat := fmt.Sprintf("%d/", i+1)
				spat := fmt.Sprintf("/%d%s", j+1, suite.IDSep)
				for _, d := range data {
					// fmt.Printf("ID = %s   %s  %s\n", d.ID, ppat, spat)
					if strings.HasPrefix(d.ID, ppat) &&
						strings.Contains(d.ID, spat) {
						fd = append(fd, d)
					}
				}
				st := statsFor(fd)
				fmt.Printf("Scenario %d %q, Test %d %q:\n",
					i+1, s.Name,
					j+1, t.File.Name,
				)
				printStat(st)
			}
		}
	*/

}

func printStat(headline string, st sdata) {
	fmt.Printf("%s Status:   Total=%d  Pass=%d (%.1f%%), Fail=%d (%.1f%%), Error=%d (%.1f%%), Bogus=%d (%.1f%%)\n",
		headline, st.n,
		st.good, 100*float64(st.good)/float64(st.n),
		st.fail, 100*float64(st.fail)/float64(st.n),
		st.erred, 100*float64(st.erred)/float64(st.n),
		st.bogus, 100*float64(st.bogus)/float64(st.n),
	)
	fmt.Printf("%s Duration: 0%%=%.1fms, 25%%=%.1fms, 50%%=%.1fms, 75%%=%.1fms, 90%%=%.1fms, 95%%=%.1fms, 99%%=%.1fms, 100%%=%.1fms\n",
		headline,
		float64(st.min/1000)/1000,
		float64(st.q25/1000)/1000,
		float64(st.median/1000)/1000,
		float64(st.q75/1000)/1000,
		float64(st.q90/1000)/1000,
		float64(st.q95/1000)/1000,
		float64(st.q99/1000)/1000,
		float64(st.max/1000)/1000,
	)
}

type sdata struct {
	n                       int
	fail, erred, bogus      int
	min, mean, max, median  time.Duration
	q25, q75, q90, q95, q99 time.Duration
	bad, good               int
	data                    []uint32 // in musec
}

func statsFor(data []bender.TestData) sdata {
	sd := sdata{
		n:    len(data),
		min:  999 * time.Hour,
		mean: time.Duration(0),
		max:  time.Duration(0),
	}
	if sd.n == 0 {
		return sd
	}

	x := make([]time.Duration, len(data))
	musec := make([]uint32, len(data))
	sum := time.Duration(0)
	for i, d := range data {
		if d.Status == ht.Fail {
			sd.fail++
		} else if d.Status == ht.Error {
			sd.erred++
		} else if d.Status == ht.Bogus {
			sd.bogus++
		}
		x[i] = d.ReqDuration
		sum += d.ReqDuration
		if d.ReqDuration/1000000 <= math.MaxUint32 {
			musec[i] = uint32(d.ReqDuration / 1000000)
		} else {
			musec[i] = math.MaxUint32
		}
	}
	sd.bad = sd.fail + sd.erred + sd.bogus
	sd.good = sd.n - sd.bad

	sd.mean = sum / time.Duration(len(data))
	sort.Sort(durationSlice(x))
	sd.min, sd.max = x[0], x[len(x)-1]
	sd.median = quantile(x, 0.5)

	sd.q25, sd.median = quantile(x, 0.25), quantile(x, 0.5)
	sd.q75, sd.q90 = quantile(x, 0.75), quantile(x, 0.90)
	sd.q95, sd.q99 = quantile(x, 0.95), quantile(x, 0.99)
	sd.data = musec
	return sd
}

type durationSlice []time.Duration

func (p durationSlice) Len() int           { return len(p) }
func (p durationSlice) Less(i, j int) bool { return p[i] < p[j] }
func (p durationSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func quantile(x []time.Duration, p float64) time.Duration {
	N := float64(len(x))
	if N == 0 {
		return 0
	} else if N == 1 {
		return x[0]
	}

	if p < 2.0/(3.0*(N+1.0/3.0)) {
		return x[0]
	}
	if p >= (N-1.0/3.0)/(N+1.0/3.0) {
		return x[len(x)-1]
	}

	h := (N+1.0/3.0)*p + 1.0/3.0
	fh := math.Floor(h)
	xl := x[int(fh)-1]
	xr := x[int(fh)]

	return time.Duration(float64(xl) + (h-fh)*float64(xr-xl))
}

func interpretLTerrors(lterr error) {
	if lterr == nil {
		fmt.Println("OKAY")
		os.Exit(0)
	}

	fmt.Println("Problems running this throughpout tests:")
	if el, ok := lterr.(ht.ErrorList); ok {
		for _, msg := range el.AsStrings() {
			fmt.Println("    ", msg)
		}
	} else {
		fmt.Println("  ", lterr.Error())
	}

	fmt.Println("PROBLEMS")
	os.Exit(1)
}

func saveLoadtestData(data []bender.TestData, failures *suite.Suite) {
	err := suite.HTMLReport(outputDir, failures)
	if err != nil {
		log.Panic(err)
	}

	file, err := os.Create(outputDir + "/throughput.csv")
	if err != nil {
		log.Panic(err)
	}
	defer file.Close()

	suite.DataToCSV(data, file)
	script := `
library(ggplot2)
d <- read.csv("throughput.csv")
d$Status <- factor(d$Status, levels <- c("NotRun", "Skipped", "Pass", "Fail", "Error", "Bogus"))

shift <- function(x, lag) {
    n <- length(x)
    xnew <- rep(NA, n)
    xnew[(lag+1):n] <- x[1:(n-lag)]
    return(xnew)
}

p <- ggplot(d, aes(x=Elapsed, y=Rate))
p <- p + geom_point(size=3) + geom_smooth()
ggsave("rate.png", plot=p, width=10, height=8, dpi=100)

d$Delta <- c(d$Elapsed[2:length(d$Elapsed)] - d$Elapsed[1:length(d$Elapsed)-1], NA)
p <- ggplot(d, aes(x=Elapsed, y=Delta)) + geom_point(size=2) + xlab("Elapsed [ms]")
ggsave("delta.png", plot=p, width=10, height=8, dpi=100)

myColors <- c("#999999", "#ffff00", "#339900", "#660000", "#ff0000", "#ff3399")
names(myColors) <- levels(d$Status)
colScale <- scale_colour_manual(name = "status",values = myColors)
fillScale <- scale_fill_manual(name = "status",values = myColors)

p <- ggplot(d, aes(x=Elapsed, y=ReqDuration, colour=Test))
p <- p + geom_point(size=3) + xlab("Elapsed [ms]") + ylab("Request Duration [ms]")
ggsave("scatter.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=Elapsed, y=ReqDuration, colour=Status))
p <- p + geom_point(size=3) + xlab("Elapsed [ms]")  + ylab("Request Duration [ms]")
p <- p + colScale
ggsave("status.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=Elapsed, y=ConcTot, colour=Status)) + colScale
p <- p + geom_point(size=3) + xlab("Elapsed [ms]") + ylab("Total Concurrency")
ggsave("conctot.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=Elapsed, y=ConcOwn, colour=Status)) + colScale
p <- p + facet_grid(Test ~ ., scales="free_y")
p <- p + geom_point(size=3) + xlab("Elapsed [ms]")
ggsave("concown.png", plot=p, width=10, height=8, dpi=100)

p <- ggplot(d, aes(x=ReqDuration, fill=Status))
p <- p + geom_histogram(binwidth=3)
p <- p + facet_grid(Test ~ ., scales="free_y")
p <- p + fillScale
ggsave("hist.png", plot=p, width=10, height=8, dpi=100)

`
	ioutil.WriteFile(outputDir+"/throughput.R", []byte(script), 0666)

}
<|endoftext|>


<commit_before>// Copyright 2017 Pilosa Corp.
//
// 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 cmd

import (
	"context"
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"github.com/pilosa/pilosa/ctl"
)

var Inspector *ctl.InspectCommand

func NewInspectCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command {
	Inspector = ctl.NewInspectCommand(os.Stdin, os.Stdout, os.Stderr)

	inspectCmd := &cobra.Command{
		Use:   "inspect",
		Short: "Get stats on a pilosa data file.",
		Long: `
Inspects a data file and provides stats.
`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 0 {
				return fmt.Errorf("path required")
			} else if len(args) > 1 {
				return fmt.Errorf("only one path allowed")
			}
			Inspector.Path = args[0]
			if err := Inspector.Run(context.Background()); err != nil {
				return err
			}
			return nil
		},
	}
	return inspectCmd
}

func init() {
	subcommandFns["inspect"] = NewInspectCommand
}
<commit_msg>Unexport cmd.Inspector<commit_after>// Copyright 2017 Pilosa Corp.
//
// 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 cmd

import (
	"context"
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"github.com/pilosa/pilosa/ctl"
)

var inspector *ctl.InspectCommand

func NewInspectCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command {
	inspector = ctl.NewInspectCommand(os.Stdin, os.Stdout, os.Stderr)

	inspectCmd := &cobra.Command{
		Use:   "inspect",
		Short: "Get stats on a pilosa data file.",
		Long: `
Inspects a data file and provides stats.
`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 0 {
				return fmt.Errorf("path required")
			} else if len(args) > 1 {
				return fmt.Errorf("only one path allowed")
			}
			inspector.Path = args[0]
			if err := inspector.Run(context.Background()); err != nil {
				return err
			}
			return nil
		},
	}
	return inspectCmd
}

func init() {
	subcommandFns["inspect"] = NewInspectCommand
}
<|endoftext|>