<commit_before>package command

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/termui"
	"github.com/golang/snappy"
	"github.com/nsf/termbox-go"
	"github.com/pmylund/sortutil"
)

var snappyMagic = []byte{130, 83, 78, 65, 80, 80, 89, 0} // SNAPPY

// TODO calculate how much data produced each day "github.com/hashicorp/go-memdb"
type Segment struct {
	Ui  cli.Ui
	Cmd string

	rootPath string
	render   bool
	filename string
	limit    int
}

func (this *Segment) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("segment", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.rootPath, "s", "", "")
	cmdFlags.BoolVar(&this.render, "render", true, "")
	cmdFlags.IntVar(&this.limit, "n", -1, "")
	cmdFlags.StringVar(&this.filename, "f", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.rootPath != "" {
		this.printSummary()
		return
	}

	if validateArgs(this, this.Ui).
		require("-f").
		invalid(args) {
		return 2
	}

	this.readSegment(this.filename)

	return
}

// a kafka message:
// crc magic attributes keyLen key messageLen message
func (this *Segment) readSegment(filename string) {
	f, err := os.Open(filename) // readonly
	swallow(err)
	defer f.Close()

	const (
		maxKeySize = 10 << 10
		maxValSize = 2 << 20
	)

	var (
		buf = make([]byte, 12)
		key = make([]byte, maxKeySize)
		val = make([]byte, maxValSize)

		msgN        int64
		firstOffset uint64 = math.MaxUint64 // sentry
		endOffset   uint64
	)
	r := bufio.NewReader(f)
	for {
		_, err := r.Read(buf)
		if err != nil {
			if err == io.EOF {
				break
			} else {
				panic(err)
			}
		}

		// offset, size, crc32, magic, attrs, key-len, key-content, msg-len, msg-content
		// crc32 = crc32(magic, attrs, key-len, key-content, msg-len, msg-content)

		// offset+size 8+4
		offset := binary.BigEndian.Uint64(buf[:8])
		size := binary.BigEndian.Uint32(buf[8:12])

		// crc32+magic+attrs+keySize[key] 4+1+1+4
		r.Read(buf[0:10])

		attr := buf[5]
		keySize := binary.BigEndian.Uint32(buf[6:10])
		if keySize > 0 && keySize != math.MaxUint32 {
			_, err = r.Read(key[:keySize])
			swallow(err)
		}

		// valSize[val] 4
		_, err = r.Read(buf[:4])
		swallow(err)
		valSize := binary.BigEndian.Uint32(buf[:4])
		if valSize > 0 {
			_, err = r.Read(val[:valSize])
			swallow(err)
		}

		switch sarama.CompressionCodec(attr) {
		case sarama.CompressionNone:
			fmt.Printf("offset:%d size:%d %s\n", offset, size, string(val[:valSize]))

		case sarama.CompressionGZIP:
			reader, err := gzip.NewReader(bytes.NewReader(val[:valSize]))
			swallow(err)
			v, err := ioutil.ReadAll(reader)
			swallow(err)

			fmt.Printf("offset:%d size:%d gzip %s\n", offset, size, string(v))

		case sarama.CompressionSnappy:
			v, err := this.snappyDecode(val[:valSize])
			swallow(err)

			fmt.Printf("offset:%d size:%d snappy %s\n", offset, size, string(v))
		}

		if firstOffset == math.MaxUint64 {
			firstOffset = offset
		}
		endOffset = offset
		msgN++
	}

	fmt.Printf("Total Messages: %d, %d - %d\n", msgN, firstOffset, endOffset)
}

func (*Segment) snappyDecode(src []byte) ([]byte, error) {
	if bytes.Equal(src[:8], snappyMagic) {
		var (
			pos   = uint32(16)
			max   = uint32(len(src))
			dst   = make([]byte, 0, len(src))
			chunk []byte
			err   error
		)
		for pos < max {
			size := binary.BigEndian.Uint32(src[pos : pos+4])
			pos += 4

			chunk, err = snappy.Decode(chunk, src[pos:pos+size])
			if err != nil {
				return nil, err
			}
			pos += size
			dst = append(dst, chunk...)
		}

		return dst, nil
	}

	return snappy.Decode(nil, src)
}

func (this *Segment) isKafkaLogSegment(fn string) bool {
	if !strings.HasSuffix(fn, ".log") || len(fn) != len("00000000000000000000.log") {
		return false
	}

	parts := strings.Split(fn, ".")
	if _, err := strconv.Atoi(parts[0]); err != nil {
		return false
	}

	return true
}

func (this *Segment) printSummary() {
	segments := make(map[string]map[int]map[int]int64) // dir:day:hour:size
	err := filepath.Walk(this.rootPath, func(path string, f os.FileInfo, err error) error {
		if f == nil {
			return err
		}
		if f.IsDir() || !this.isKafkaLogSegment(f.Name()) {
			return nil
		}

		dir := filepath.Base(filepath.Dir(path))
		if _, present := segments[dir]; !present {
			segments[dir] = make(map[int]map[int]int64)
		}
		if _, present := segments[dir][f.ModTime().Day()]; !present {
			segments[dir][f.ModTime().Day()] = make(map[int]int64)
		}
		segments[dir][f.ModTime().Day()][f.ModTime().Hour()] += f.Size()
		return nil
	})
	if err != nil {
		this.Ui.Error(err.Error())
	}

	if this.render {
		this.doRender(segments)
	}

	partitions := make([]string, 0, len(segments))
	for dir := range segments {
		partitions = append(partitions, dir)
	}
	sort.Strings(partitions)

	type segment struct {
		partition string
		day       int
		hour      int
		size      int64
	}

	var maxSegment segment
	var totalSize int64
	for _, p := range partitions {
		summary := make([]segment, 0)
		for day, hourSize := range segments[p] {
			for hour, size := range hourSize {
				summary = append(summary, segment{
					partition: p,
					day:       day,
					hour:      hour,
					size:      size,
				})
			}
		}
		sortutil.AscByField(summary, "size")
		if this.limit > 0 && len(summary) > this.limit {
			summary = summary[:this.limit]
		}

		for _, s := range summary {
			if s.size > maxSegment.size {
				maxSegment = s
			}

			totalSize += s.size
			this.Ui.Output(fmt.Sprintf("%50s day:%2d hour:%2d size:%s", p,
				s.day, s.hour, gofmt.ByteSize(s.size)))
		}

	}

	this.Ui.Output(fmt.Sprintf("%50s day:%2d hour:%2d size:%s", "MAX-"+maxSegment.partition,
		maxSegment.day, maxSegment.hour, gofmt.ByteSize(maxSegment.size)))
	this.Ui.Output(fmt.Sprintf("%50s %s", "-TOTAL-", gofmt.ByteSize(totalSize)))

	return
}

func (this *Segment) doRender(segments map[string]map[int]map[int]int64) {
	err := termui.Init()
	swallow(err)
	defer termui.Close()

	if len(segments) > 1 {
		panic("recursive dir not allowed")
	}

	termui.UseTheme("helloworld")

	w, h := termbox.Size()

	bc1 := termui.NewBarChart()
	bc1.Border.Label = "Segement Size in GB"

	dayHours := make([]string, 0)
	m := make(map[string]int)
	for _, dayHourSize := range segments {
		for day, hourSize := range dayHourSize {
			for hour, size := range hourSize {
				dh := fmt.Sprintf("%d:%d", day, hour)
				dayHours = append(dayHours, dh)
				m[dh] = int(size / (1 << 30))
			}
		}
	}
	sort.Strings(dayHours)

	sizesInGB := make([]int, 0)
	hours := make([]string, 0, len(dayHours))
	for _, dh := range dayHours {
		sizesInGB = append(sizesInGB, m[dh])

		tuples := strings.Split(dh, ":")
		hours = append(hours, tuples[1])
	}

	bc1.Data = sizesInGB
	bc1.DataLabels = hours

	bc1.Width = w
	bc1.SetY(0)
	bc1.Height = h
	bc1.TextColor = termui.ColorWhite
	bc1.BarColor = termui.ColorRed
	bc1.NumColor = termui.ColorYellow

	termui.Render(bc1)
	termbox.PollEvent()
}

func (*Segment) Synopsis() string {
	return "Scan the kafka segments and display summary"
}

func (this *Segment) Help() string {
	help := fmt.Sprintf(`
Usage: %s segment [options]

    %s

    -f segment file name
     Display segment content.

    -n limit
     Default unlimited.

    -s dir
     Sumamry of a segment dir.
     Summary across partitions is supported if they have the same parent dir.

    -render
     Render histogram. Work with '-s <dir>'.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>fit in screen<commit_after>package command

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/termui"
	"github.com/golang/snappy"
	"github.com/nsf/termbox-go"
	"github.com/pmylund/sortutil"
)

var snappyMagic = []byte{130, 83, 78, 65, 80, 80, 89, 0} // SNAPPY

// TODO calculate how much data produced each day "github.com/hashicorp/go-memdb"
type Segment struct {
	Ui  cli.Ui
	Cmd string

	rootPath string
	render   bool
	filename string
	limit    int
}

func (this *Segment) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("segment", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.rootPath, "s", "", "")
	cmdFlags.BoolVar(&this.render, "render", true, "")
	cmdFlags.IntVar(&this.limit, "n", -1, "")
	cmdFlags.StringVar(&this.filename, "f", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.rootPath != "" {
		this.printSummary()
		return
	}

	if validateArgs(this, this.Ui).
		require("-f").
		invalid(args) {
		return 2
	}

	this.readSegment(this.filename)

	return
}

// a kafka message:
// crc magic attributes keyLen key messageLen message
func (this *Segment) readSegment(filename string) {
	f, err := os.Open(filename) // readonly
	swallow(err)
	defer f.Close()

	const (
		maxKeySize = 10 << 10
		maxValSize = 2 << 20
	)

	var (
		buf = make([]byte, 12)
		key = make([]byte, maxKeySize)
		val = make([]byte, maxValSize)

		msgN        int64
		firstOffset uint64 = math.MaxUint64 // sentry
		endOffset   uint64
	)
	r := bufio.NewReader(f)
	for {
		_, err := r.Read(buf)
		if err != nil {
			if err == io.EOF {
				break
			} else {
				panic(err)
			}
		}

		// offset, size, crc32, magic, attrs, key-len, key-content, msg-len, msg-content
		// crc32 = crc32(magic, attrs, key-len, key-content, msg-len, msg-content)

		// offset+size 8+4
		offset := binary.BigEndian.Uint64(buf[:8])
		size := binary.BigEndian.Uint32(buf[8:12])

		// crc32+magic+attrs+keySize[key] 4+1+1+4
		r.Read(buf[0:10])

		attr := buf[5]
		keySize := binary.BigEndian.Uint32(buf[6:10])
		if keySize > 0 && keySize != math.MaxUint32 {
			_, err = r.Read(key[:keySize])
			swallow(err)
		}

		// valSize[val] 4
		_, err = r.Read(buf[:4])
		swallow(err)
		valSize := binary.BigEndian.Uint32(buf[:4])
		if valSize > 0 {
			_, err = r.Read(val[:valSize])
			swallow(err)
		}

		switch sarama.CompressionCodec(attr) {
		case sarama.CompressionNone:
			fmt.Printf("offset:%d size:%d %s\n", offset, size, string(val[:valSize]))

		case sarama.CompressionGZIP:
			reader, err := gzip.NewReader(bytes.NewReader(val[:valSize]))
			swallow(err)
			v, err := ioutil.ReadAll(reader)
			swallow(err)

			fmt.Printf("offset:%d size:%d gzip %s\n", offset, size, string(v))

		case sarama.CompressionSnappy:
			v, err := this.snappyDecode(val[:valSize])
			swallow(err)

			fmt.Printf("offset:%d size:%d snappy %s\n", offset, size, string(v))
		}

		if firstOffset == math.MaxUint64 {
			firstOffset = offset
		}
		endOffset = offset
		msgN++
	}

	fmt.Printf("Total Messages: %d, %d - %d\n", msgN, firstOffset, endOffset)
}

func (*Segment) snappyDecode(src []byte) ([]byte, error) {
	if bytes.Equal(src[:8], snappyMagic) {
		var (
			pos   = uint32(16)
			max   = uint32(len(src))
			dst   = make([]byte, 0, len(src))
			chunk []byte
			err   error
		)
		for pos < max {
			size := binary.BigEndian.Uint32(src[pos : pos+4])
			pos += 4

			chunk, err = snappy.Decode(chunk, src[pos:pos+size])
			if err != nil {
				return nil, err
			}
			pos += size
			dst = append(dst, chunk...)
		}

		return dst, nil
	}

	return snappy.Decode(nil, src)
}

func (this *Segment) isKafkaLogSegment(fn string) bool {
	if !strings.HasSuffix(fn, ".log") || len(fn) != len("00000000000000000000.log") {
		return false
	}

	parts := strings.Split(fn, ".")
	if _, err := strconv.Atoi(parts[0]); err != nil {
		return false
	}

	return true
}

func (this *Segment) printSummary() {
	segments := make(map[string]map[int]map[int]int64) // dir:day:hour:size
	err := filepath.Walk(this.rootPath, func(path string, f os.FileInfo, err error) error {
		if f == nil {
			return err
		}
		if f.IsDir() || !this.isKafkaLogSegment(f.Name()) {
			return nil
		}

		dir := filepath.Base(filepath.Dir(path))
		if _, present := segments[dir]; !present {
			segments[dir] = make(map[int]map[int]int64)
		}
		if _, present := segments[dir][f.ModTime().Day()]; !present {
			segments[dir][f.ModTime().Day()] = make(map[int]int64)
		}
		segments[dir][f.ModTime().Day()][f.ModTime().Hour()] += f.Size()
		return nil
	})
	if err != nil {
		this.Ui.Error(err.Error())
	}

	if this.render {
		this.doRender(segments)
	}

	partitions := make([]string, 0, len(segments))
	for dir := range segments {
		partitions = append(partitions, dir)
	}
	sort.Strings(partitions)

	type segment struct {
		partition string
		day       int
		hour      int
		size      int64
	}

	var maxSegment segment
	var totalSize int64
	for _, p := range partitions {
		summary := make([]segment, 0)
		for day, hourSize := range segments[p] {
			for hour, size := range hourSize {
				summary = append(summary, segment{
					partition: p,
					day:       day,
					hour:      hour,
					size:      size,
				})
			}
		}
		sortutil.AscByField(summary, "size")
		if this.limit > 0 && len(summary) > this.limit {
			summary = summary[:this.limit]
		}

		for _, s := range summary {
			if s.size > maxSegment.size {
				maxSegment = s
			}

			totalSize += s.size
			this.Ui.Output(fmt.Sprintf("%50s day:%2d hour:%2d size:%s", p,
				s.day, s.hour, gofmt.ByteSize(s.size)))
		}

	}

	this.Ui.Output(fmt.Sprintf("%50s day:%2d hour:%2d size:%s", "MAX-"+maxSegment.partition,
		maxSegment.day, maxSegment.hour, gofmt.ByteSize(maxSegment.size)))
	this.Ui.Output(fmt.Sprintf("%50s %s", "-TOTAL-", gofmt.ByteSize(totalSize)))

	return
}

func (this *Segment) doRender(segments map[string]map[int]map[int]int64) {
	err := termui.Init()
	swallow(err)
	defer termui.Close()

	if len(segments) > 1 {
		panic("recursive dir not allowed")
	}

	termui.UseTheme("helloworld")

	w, h := termbox.Size()

	bc1 := termui.NewBarChart()
	bc1.Border.Label = "Segement Size in GB"

	dayHours := make([]string, 0)
	m := make(map[string]int)
	for _, dayHourSize := range segments {
		for day, hourSize := range dayHourSize {
			for hour, size := range hourSize {
				dh := fmt.Sprintf("%d:%d", day, hour)
				dayHours = append(dayHours, dh)
				m[dh] = int(size / (1 << 30))
			}
		}
	}
	sort.Strings(dayHours)

	sizesInGB := make([]int, 0)
	hours := make([]string, 0, len(dayHours))
	for _, dh := range dayHours {
		sizesInGB = append(sizesInGB, m[dh])

		tuples := strings.Split(dh, ":")
		hours = append(hours, tuples[1])
	}

	maxItems := (w - 2) / 4
	begin := 0
	if len(sizesInGB) > maxItems {
		begin = len(sizesInGB) - maxItems
	}

	bc1.Data = sizesInGB[begin:]
	bc1.DataLabels = hours[begin:]

	bc1.Width = w
	bc1.SetY(0)
	bc1.Height = h
	bc1.TextColor = termui.ColorWhite
	bc1.BarColor = termui.ColorRed
	bc1.NumColor = termui.ColorYellow

	termui.Render(bc1)
	termbox.PollEvent()
}

func (*Segment) Synopsis() string {
	return "Scan the kafka segments and display summary"
}

func (this *Segment) Help() string {
	help := fmt.Sprintf(`
Usage: %s segment [options]

    %s

    -f segment file name
     Display segment content.

    -n limit
     Default unlimited.

    -s dir
     Sumamry of a segment dir.
     Summary across partitions is supported if they have the same parent dir.

    -render
     Render histogram. Work with '-s <dir>'.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/environs"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju-core/juju/environs/dummy"
	_ "launchpad.net/juju-core/juju/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	// TODO(dfc) place the logic in a loop with a suitable delay
	p, err := NewProvisioner(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	return p.Wait()
}

type Provisioner struct {
	st      *state.State
	info    *state.Info
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher

	// machine.Id => environs.Instance
	instances map[int]environs.Instance
	// instance.Id => *state.Machine
	machines map[string]*state.Machine
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(info *state.Info) (*Provisioner, error) {
	st, err := state.Open(info)
	if err != nil {
		return nil, err
	}
	p := &Provisioner{
		st:        st,
		info:      info,
		instances: make(map[int]environs.Instance),
		machines:  make(map[string]*state.Machine),
	}
	go p.loop()
	return p, nil
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	defer p.st.Close()
	p.environWatcher = p.st.WatchEnvironConfig()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			log.Printf("provisioner loaded new environment configuration")

			// call processMachines to stop any unknown instances before watching machines.
			if err := p.processMachines(nil, nil); err != nil {
				p.tomb.Kill(err)
			}

			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	p.machinesWatcher = p.st.WatchMachines()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner loaded new environment configuration")
		case machines, ok := <-p.machinesWatcher.Changes():
			if !ok {
				err := p.machinesWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			// TODO(dfc) fire process machines periodically to shut down unknown
			// instances.
			if err := p.processMachines(machines.Added, machines.Deleted); err != nil {
				p.tomb.Kill(err)
			}
		}
	}
}

// Wait waits for the Provisioner to exit.
func (p *Provisioner) Wait() error {
	return p.tomb.Wait()
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

func (p *Provisioner) processMachines(added, removed []*state.Machine) error {
	// step 1. find which of the added machines have not
	// yet been allocated a started instance.
	notstarted, err := p.findNotStarted(added)
	if err != nil {
		return err
	}

	// step 2. start an instance for any machines we found.
	if err := p.startMachines(notstarted); err != nil {
		return err
	}

	// step 3. stop all machines that were removed from the state.
	stopping, err := p.instancesForMachines(removed)
	if err != nil {
		return err
	}

	// step 4. find instances which are running but have no machine
	// associated with them.
	unknown, err := p.findUnknownInstances()

	return p.stopInstances(append(stopping, unknown...))
}

// findUnknownInstances finds instances which are not associated with a machine.
func (p *Provisioner) findUnknownInstances() ([]environs.Instance, error) {
	all, err := p.environ.AllInstances()
	if err != nil {
		return nil, err
	}
	instances := make(map[string]environs.Instance)
	for _, i := range all {
		instances[i.Id()] = i
	}
	// TODO(dfc) this is very inefficient, p.machines cache may help.
	machines, err := p.st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			continue
		}
		delete(instances, id)
	}
	var unknown []environs.Instance
	for _, i := range instances {
		unknown = append(unknown, i)
	}
	return unknown, nil
}

// findNotStarted finds machines without an InstanceId set, these are defined as not started.
func (p *Provisioner) findNotStarted(machines []*state.Machine) ([]*state.Machine, error) {
	var notstarted []*state.Machine
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			notstarted = append(notstarted, m)
		} else {
			log.Printf("machine %s already started as instance %q", m, id)
		}
	}
	return notstarted, nil
}

func (p *Provisioner) startMachines(machines []*state.Machine) error {
	for _, m := range machines {
		if err := p.startMachine(m); err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) startMachine(m *state.Machine) error {
	// TODO(dfc) the state.Info passed to environ.StartInstance remains contentious
	// however as the PA only knows one state.Info, and that info is used by MAs and
	// UAs to locate the ZK for this environment, it is logical to use the same
	// state.Info as the PA.
	inst, err := p.environ.StartInstance(m.Id(), p.info)
	if err != nil {
		log.Printf("provisioner can't start machine %s: %v", m, err)
		return err
	}

	// assign the instance id to the machine
	if err := m.SetInstanceId(inst.Id()); err != nil {
		return err
	}

	// populate the local cache
	p.instances[m.Id()] = inst
	p.machines[inst.Id()] = m
	log.Printf("provisioner started machine %s as instance %s", m, inst.Id())
	return nil
}

func (p *Provisioner) stopInstances(instances []environs.Instance) error {
	// Although calling StopInstance with an empty slice should produce no change in the
	// provider, environs like dummy do not consider this a noop.
	if len(instances) == 0 {
		return nil
	}
	if err := p.environ.StopInstances(instances); err != nil {
		return err
	}

	// cleanup cache
	for _, i := range instances {
		if m, ok := p.machines[i.Id()]; ok {
			delete(p.machines, i.Id())
			delete(p.instances, m.Id())
		}
	}
	return nil
}

// instanceForMachine returns the environs.Instance that represents this machine's instance.
func (p *Provisioner) instanceForMachine(m *state.Machine) (environs.Instance, error) {
	inst, ok := p.instances[m.Id()]
	if !ok {
		// not cached locally, ask the environ.
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			return nil, fmt.Errorf("machine %s not found", m)
		}
		// TODO(dfc) this should be batched, or the cache preloaded at startup to
		// avoid N calls to the envirion.
		insts, err := p.environ.Instances([]string{id})
		if err != nil {
			// the provider doesn't know about this instance, give up.
			return nil, err
		}
		inst = insts[0]
	}
	return inst, nil
}

// instancesForMachines returns a list of environs.Instance that represent the list of machines running
// in the provider.
func (p *Provisioner) instancesForMachines(machines []*state.Machine) ([]environs.Instance, error) {
	var insts []environs.Instance
	for _, m := range machines {
		inst, err := p.instanceForMachine(m)
		if err != nil {
			return nil, err
		}
		insts = append(insts, inst)
	}
	return insts, nil
}
<commit_msg>responding to review feedback<commit_after>package main

import (
	"fmt"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/environs"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju-core/juju/environs/dummy"
	_ "launchpad.net/juju-core/juju/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	// TODO(dfc) place the logic in a loop with a suitable delay
	p, err := NewProvisioner(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	return p.Wait()
}

type Provisioner struct {
	st      *state.State
	info    *state.Info
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher

	// machine.Id => environs.Instance
	instances map[int]environs.Instance
	// instance.Id => *state.Machine
	machines map[string]*state.Machine
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(info *state.Info) (*Provisioner, error) {
	st, err := state.Open(info)
	if err != nil {
		return nil, err
	}
	p := &Provisioner{
		st:        st,
		info:      info,
		instances: make(map[int]environs.Instance),
		machines:  make(map[string]*state.Machine),
	}
	go p.loop()
	return p, nil
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	defer p.st.Close()
	p.environWatcher = p.st.WatchEnvironConfig()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			log.Printf("provisioner loaded new environment configuration")

			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// call processMachines to stop any unknown instances before watching machines.
	if err := p.processMachines(nil, nil); err != nil {
		p.tomb.Kill(err)
	}
	p.machinesWatcher = p.st.WatchMachines()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner loaded new environment configuration")
		case machines, ok := <-p.machinesWatcher.Changes():
			if !ok {
				err := p.machinesWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			// TODO(dfc) fire process machines periodically to shut down unknown
			// instances.
			if err := p.processMachines(machines.Added, machines.Deleted); err != nil {
				p.tomb.Kill(err)
			}
		}
	}
}

// Wait waits for the Provisioner to exit.
func (p *Provisioner) Wait() error {
	return p.tomb.Wait()
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

func (p *Provisioner) processMachines(added, removed []*state.Machine) error {
	// step 1. find which of the added machines have not
	// yet been allocated a started instance.
	notstarted, err := p.findNotStarted(added)
	if err != nil {
		return err
	}

	// step 2. start an instance for any machines we found.
	if err := p.startMachines(notstarted); err != nil {
		return err
	}

	// step 3. stop all machines that were removed from the state.
	stopping, err := p.instancesForMachines(removed)
	if err != nil {
		return err
	}

	// step 4. find instances which are running but have no machine
	// associated with them.
	unknown, err := p.findUnknownInstances()

	return p.stopInstances(append(stopping, unknown...))
}

// findUnknownInstances finds instances which are not associated with a machine.
func (p *Provisioner) findUnknownInstances() ([]environs.Instance, error) {
	all, err := p.environ.AllInstances()
	if err != nil {
		return nil, err
	}
	instances := make(map[string]environs.Instance)
	for _, i := range all {
		instances[i.Id()] = i
	}
	// TODO(dfc) this is very inefficient, p.machines cache may help.
	machines, err := p.st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			continue
		}
		delete(instances, id)
	}
	var unknown []environs.Instance
	for _, i := range instances {
		unknown = append(unknown, i)
	}
	return unknown, nil
}

// findNotStarted finds machines without an InstanceId set, these are defined as not started.
func (p *Provisioner) findNotStarted(machines []*state.Machine) ([]*state.Machine, error) {
	var notstarted []*state.Machine
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			notstarted = append(notstarted, m)
		} else {
			log.Printf("machine %s already started as instance %q", m, id)
		}
	}
	return notstarted, nil
}

func (p *Provisioner) startMachines(machines []*state.Machine) error {
	for _, m := range machines {
		if err := p.startMachine(m); err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) startMachine(m *state.Machine) error {
	// TODO(dfc) the state.Info passed to environ.StartInstance remains contentious
	// however as the PA only knows one state.Info, and that info is used by MAs and
	// UAs to locate the ZK for this environment, it is logical to use the same
	// state.Info as the PA.
	inst, err := p.environ.StartInstance(m.Id(), p.info)
	if err != nil {
		log.Printf("provisioner can't start machine %s: %v", m, err)
		return err
	}

	// assign the instance id to the machine
	if err := m.SetInstanceId(inst.Id()); err != nil {
		return err
	}

	// populate the local cache
	p.instances[m.Id()] = inst
	p.machines[inst.Id()] = m
	log.Printf("provisioner started machine %s as instance %s", m, inst.Id())
	return nil
}

func (p *Provisioner) stopInstances(instances []environs.Instance) error {
	// Although calling StopInstance with an empty slice should produce no change in the
	// provider, environs like dummy do not consider this a noop.
	if len(instances) == 0 {
		return nil
	}
	if err := p.environ.StopInstances(instances); err != nil {
		return err
	}

	// cleanup cache
	for _, i := range instances {
		if m, ok := p.machines[i.Id()]; ok {
			delete(p.machines, i.Id())
			delete(p.instances, m.Id())
		}
	}
	return nil
}

// instanceForMachine returns the environs.Instance that represents this machine's instance.
func (p *Provisioner) instanceForMachine(m *state.Machine) (environs.Instance, error) {
	inst, ok := p.instances[m.Id()]
	if !ok {
		// not cached locally, ask the environ.
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			return nil, fmt.Errorf("machine %s not found", m)
		}
		// TODO(dfc) this should be batched, or the cache preloaded at startup to
		// avoid N calls to the envirion.
		insts, err := p.environ.Instances([]string{id})
		if err != nil {
			// the provider doesn't know about this instance, give up.
			return nil, err
		}
		inst = insts[0]
	}
	return inst, nil
}

// instancesForMachines returns a list of environs.Instance that represent the list of machines running
// in the provider.
func (p *Provisioner) instancesForMachines(machines []*state.Machine) ([]environs.Instance, error) {
	var insts []environs.Instance
	for _, m := range machines {
		inst, err := p.instanceForMachine(m)
		if err != nil {
			return nil, err
		}
		insts = append(insts, inst)
	}
	return insts, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/siddontang/go/num"
	"github.com/siddontang/ledisdb/config"
	"github.com/siddontang/ledisdb/ledis"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

var KB = config.KB
var MB = config.MB
var GB = config.GB

var name = flag.String("db_name", "goleveldb", "db name")
var number = flag.Int("n", 10000, "request number")
var clients = flag.Int("c", 50, "number of clients")
var round = flag.Int("r", 1, "benchmark round number")
var valueSize = flag.Int("vsize", 100, "kv value size")
var wg sync.WaitGroup

var ldb *ledis.Ledis
var db *ledis.DB

var loop int = 0

func bench(cmd string, f func()) {
	wg.Add(*clients)

	t1 := time.Now()
	for i := 0; i < *clients; i++ {
		go func() {
			for j := 0; j < loop; j++ {
				f()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	t2 := time.Now()

	d := t2.Sub(t1)
	fmt.Printf("%s %s: %0.3f micros/op, %0.2fmb/s %0.2fop/s\n",
		cmd,
		d.String(),
		float64(d.Nanoseconds()/1e3)/float64(*number),
		float64((*valueSize+16)*(*number))/(1024.0*1024.0*(d.Seconds())),
		float64(*number)/d.Seconds())
}

var kvSetBase int64 = 0
var kvGetBase int64 = 0

func benchSet() {
	f := func() {
		value := make([]byte, *valueSize)
		n := atomic.AddInt64(&kvSetBase, 1)

		db.Set(num.Int64ToBytes(n), value)
	}

	bench("set", f)
}

func benchGet() {
	f := func() {
		n := atomic.AddInt64(&kvGetBase, 1)
		v, err := db.Get(num.Int64ToBytes(n))
		if err != nil {
			println(err.Error())
		} else if len(v) != *valueSize {
			println(len(v), *valueSize)
		}
	}

	bench("get", f)
}

func setRocksDB(cfg *config.RocksDBConfig) {
	cfg.BlockSize = 64 * KB
	cfg.WriteBufferSize = 64 * MB
	cfg.MaxWriteBufferNum = 2
	cfg.MaxBytesForLevelBase = 512 * MB
	cfg.TargetFileSizeBase = 64 * MB
	cfg.BackgroundThreads = 4
	cfg.HighPriorityBackgroundThreads = 1
	cfg.MaxBackgroundCompactions = 3
	cfg.MaxBackgroundFlushes = 1
	cfg.CacheSize = 512 * MB
	cfg.EnableStatistics = true
	cfg.StatsDumpPeriodSec = 5
	cfg.Level0FileNumCompactionTrigger = 8
	cfg.MaxBytesForLevelMultiplier = 8
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	cfg := config.NewConfigDefault()
	cfg.DataDir = "./var/ledis_dbbench"
	cfg.DBName = *name
	os.RemoveAll(cfg.DBPath)
	defer os.RemoveAll(cfg.DBPath)

	os.MkdirAll(cfg.DBPath, 0755)

	cfg.LevelDB.BlockSize = 32 * KB
	cfg.LevelDB.CacheSize = 512 * MB
	cfg.LevelDB.WriteBufferSize = 64 * MB
	cfg.LevelDB.MaxOpenFiles = 1000

	setRocksDB(&cfg.RocksDB)

	var err error
	ldb, err = ledis.Open(cfg)
	if err != nil {
		println(err.Error())
		return
	}

	db, _ = ldb.Select(0)

	if *number <= 0 {
		panic("invalid number")
		return
	}

	if *clients <= 0 || *number < *clients {
		panic("invalid client number")
		return
	}

	loop = *number / *clients

	if *round <= 0 {
		*round = 1
	}

	for i := 0; i < *round; i++ {
		benchSet()
		benchGet()

		println("")
	}
}
<commit_msg>add get slice for dbbench<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/siddontang/go/num"
	"github.com/siddontang/ledisdb/config"
	"github.com/siddontang/ledisdb/ledis"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

var KB = config.KB
var MB = config.MB
var GB = config.GB

var name = flag.String("db_name", "goleveldb", "db name")
var number = flag.Int("n", 10000, "request number")
var clients = flag.Int("c", 50, "number of clients")
var round = flag.Int("r", 1, "benchmark round number")
var valueSize = flag.Int("vsize", 100, "kv value size")
var wg sync.WaitGroup

var ldb *ledis.Ledis
var db *ledis.DB

var loop int = 0

func bench(cmd string, f func()) {
	wg.Add(*clients)

	t1 := time.Now()
	for i := 0; i < *clients; i++ {
		go func() {
			for j := 0; j < loop; j++ {
				f()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	t2 := time.Now()

	d := t2.Sub(t1)
	fmt.Printf("%s %s: %0.3f micros/op, %0.2fmb/s %0.2fop/s\n",
		cmd,
		d.String(),
		float64(d.Nanoseconds()/1e3)/float64(*number),
		float64((*valueSize+16)*(*number))/(1024.0*1024.0*(d.Seconds())),
		float64(*number)/d.Seconds())
}

var kvSetBase int64 = 0
var kvGetBase int64 = 0

func benchSet() {
	f := func() {
		value := make([]byte, *valueSize)
		n := atomic.AddInt64(&kvSetBase, 1)

		db.Set(num.Int64ToBytes(n), value)
	}

	bench("set", f)
}

func benchGet() {
	kvGetBase = 0
	f := func() {
		n := atomic.AddInt64(&kvGetBase, 1)
		v, err := db.Get(num.Int64ToBytes(n))
		if err != nil {
			println(err.Error())
		} else if len(v) != *valueSize {
			println(len(v), *valueSize)
		}
	}

	bench("get", f)
}

var kvGetSliceBase int64 = 0

func benchGetSlice() {
	kvGetSliceBase = 0
	f := func() {
		n := atomic.AddInt64(&kvGetSliceBase, 1)
		v, err := db.GetSlice(num.Int64ToBytes(n))
		if err != nil {
			println(err.Error())
		} else if v != nil {
			v.Free()
		}
	}

	bench("getslice", f)
}

func setRocksDB(cfg *config.RocksDBConfig) {
	cfg.BlockSize = 64 * KB
	cfg.WriteBufferSize = 64 * MB
	cfg.MaxWriteBufferNum = 2
	cfg.MaxBytesForLevelBase = 512 * MB
	cfg.TargetFileSizeBase = 64 * MB
	cfg.BackgroundThreads = 4
	cfg.HighPriorityBackgroundThreads = 1
	cfg.MaxBackgroundCompactions = 3
	cfg.MaxBackgroundFlushes = 1
	cfg.CacheSize = 512 * MB
	cfg.EnableStatistics = true
	cfg.StatsDumpPeriodSec = 5
	cfg.Level0FileNumCompactionTrigger = 8
	cfg.MaxBytesForLevelMultiplier = 8
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	cfg := config.NewConfigDefault()
	cfg.DataDir = "./var/ledis_dbbench"
	cfg.DBName = *name
	os.RemoveAll(cfg.DBPath)
	defer os.RemoveAll(cfg.DBPath)

	os.MkdirAll(cfg.DBPath, 0755)

	cfg.LevelDB.BlockSize = 32 * KB
	cfg.LevelDB.CacheSize = 512 * MB
	cfg.LevelDB.WriteBufferSize = 64 * MB
	cfg.LevelDB.MaxOpenFiles = 1000

	setRocksDB(&cfg.RocksDB)

	var err error
	ldb, err = ledis.Open(cfg)
	if err != nil {
		println(err.Error())
		return
	}

	db, _ = ldb.Select(0)

	if *number <= 0 {
		panic("invalid number")
		return
	}

	if *clients <= 0 || *number < *clients {
		panic("invalid client number")
		return
	}

	loop = *number / *clients

	if *round <= 0 {
		*round = 1
	}

	for i := 0; i < *round; i++ {
		benchSet()
		benchGet()
		benchGetSlice()
		benchGet()
		benchGetSlice()

		println("")
	}
}
<|endoftext|>


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

// THIS FILE IS DEPRECATED!!!
// Please edit the new "jiri" tool in release/go/src/v.io/jiri.

package main

import (
	"os"
	"os/exec"
)

// main calls "jiri" tool with whatever arguments it was called with.
func main() {
	args := os.Args[1:]
	cmd := exec.Command("jiri", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	// If JIRI_ROOT is not set, set it to $V23_ROOT.
	jiriRoot := os.Getenv("JIRI_ROOT")
	if jiriRoot == "" {
		if err := os.Setenv("JIRI_ROOT", os.Getenv("V23_ROOT")); err != nil {
			panic(err)
		}
	}

	if err := cmd.Run(); err != nil {
		// The jiri tool should have reported an error in its output.  Don't
		// print an error here because it can be confusing and makes it harder
		// to spot the real error.
		os.Exit(1)
	}
}
<commit_msg>devtools/v23: Add deprecation notice for v23 tool.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// THIS FILE IS DEPRECATED!!!
// Please edit the new "jiri" tool in release/go/src/v.io/jiri.

package main

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

// main calls "jiri" tool with whatever arguments it was called with.
func main() {
	args := os.Args[1:]
	cmd := exec.Command("jiri", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	// If JIRI_ROOT is not set, set it to $V23_ROOT.
	jiriRoot := os.Getenv("JIRI_ROOT")
	if jiriRoot == "" {
		if err := os.Setenv("JIRI_ROOT", os.Getenv("V23_ROOT")); err != nil {
			panic(err)
		}
	}

	fmt.Fprintf(os.Stderr, "\nWARNING: The v23 tool will soon be deprecated.\nPlease run 'jiri %s' instead.\n\n", strings.Join(args, " "))

	// Sleep for annoyance.
	time.Sleep(3 * time.Second)

	if err := cmd.Run(); err != nil {
		// The jiri tool should have reported an error in its output.  Don't
		// print an error here because it can be confusing and makes it harder
		// to spot the real error.
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"math"
	"net/http"
	"strconv"
	"time"

	"github.com/codahale/hdrhistogram"

	"golang.org/x/image/font"
	"golang.org/x/image/font/inconsolata"
	"golang.org/x/image/math/fixed"
)

// TODO(kr): collect more fine-grain stats

const (
	// Use these to go to p99.99999
	// pixel0Quantile = 0.5
	// decayPerPixel  = 0.98
	pixel0Quantile = 0.4
	decayPerPixel  = 0.99
)

var (
	dims = image.Rect(0, 0, 800, 100)

	// In reverse order, newest to oldest,
	// so strongest to faintest color.
	// Don't forget, color.RGBA fields are
	// alpha premultiplied.
	lineColors = [...]color.Color{
		color.RGBA{R: 0xff, A: 0xff},
		color.RGBA{R: 0xcc, A: 0xcc},
		color.RGBA{R: 0x88, A: 0x88},
		color.RGBA{R: 0x55, A: 0x55},
		color.RGBA{R: 0x22, A: 0x22},
	}
	curLineColor = color.RGBA{A: 0x55}
	labelColor   = color.Black
	drawer       = font.Drawer{
		Src:  image.NewUniform(labelColor),
		Face: inconsolata.Regular8x16,
	}
)

func heatmap(w http.ResponseWriter, req *http.Request) {
	query := req.URL.Query()
	name := query.Get("name")
	id, err := strconv.Atoi(query.Get("id"))
	if err != nil {
		pngError(w, err)
		return
	}

	dv := getDebugVars(id)
	if dv == nil {
		pngError(w, errors.New("not found"))
		return
	}

	latency, ok := dv.Latency[name]
	if !ok {
		pngError(w, errors.New("not found"))
		return
	}

	img := newImage(dims)
	d := drawer
	d.Dst = img

	var (
		hists   []*hdrhistogram.Histogram
		trueMax time.Duration
		over    int
	)
	for _, b := range latency.Buckets {
		hists = append(hists, hdrhistogram.Import(&b.Histogram))
		over += b.Over
		if d := time.Duration(b.Max); d > trueMax {
			trueMax = d
		}
	}

	var max int64
	for _, hist := range hists {
		if v := hist.Max(); v > max {
			max = v
		}
	}

	if max == 0 {
		drawf(d, 4, 20, "no histogram data")
	} else {
		max += max / 10
		max = roundms(max)
		if max < int64(10*time.Millisecond) {
			max += int64(time.Millisecond)
		}
		complete := hists[:len(hists)-1]
		for i, hist := range complete {
			rindex := len(complete) - i - 1
			graph(img, hist, max, lineColors[rindex%len(lineColors)])
		}
		// special color for incomplete bucket
		graph(img, hists[len(hists)-1], max, curLineColor)
		drawf(d, 4, 20, "%v", time.Duration(max))
	}
	drawf(d, 4, 38, "%s", name)
	drawf(d, 4, 54, "over: %d   max: %v", over, trueMax)
	label(img, labelColor)
	png.Encode(w, img)
}

func graph(img *image.RGBA, hist *hdrhistogram.Histogram, ymax int64, color color.Color) {
	graph := img.SubImage(img.Bounds().Inset(2)).(*image.RGBA)
	gdims := graph.Bounds()
	d := drawer
	d.Dst = graph

	labelPixels := 50
	prevY := int(scale(valueAtPixel(hist, 0), ymax, int64(gdims.Dy())))
	for x := 1; x < gdims.Dx(); x++ {
		v := valueAtPixel(hist, x)
		y := int(scale(v, ymax, int64(gdims.Dy())))
		vLineSeg(graph, gdims.Min.X+x, gdims.Max.Y-y-1, gdims.Max.Y-prevY-1, color)
		labelPixels++
		prevY = y
	}
}

func vLineSeg(img *image.RGBA, x, y0, y1 int, color color.Color) {
	dims := image.Rect(x, y0, x+1, y1+1)
	draw.Draw(img, dims, image.NewUniform(color), image.ZP, draw.Over)
}

func label(img *image.RGBA, color color.Color) {
	graph := img.SubImage(img.Bounds().Inset(2)).(*image.RGBA)
	gdims := graph.Bounds()
	d := drawer
	d.Dst = graph

	labelDigits := 0
	labelPixels := 0
	for x := 0; x < gdims.Dx(); x++ {
		q := quantileAtPixel(x)
		if dig := digits(1 - q); labelPixels >= 50 && dig > labelDigits {
			labelPixels = 0
			labelDigits = dig
			for i := 0; i < 5; i++ {
				graph.Set(gdims.Min.X+x, gdims.Max.Y-i-1, color)
			}
			prec := labelDigits - 3
			if prec < 0 {
				prec = 0
			}
			drawf(d, gdims.Min.X+x+2, gdims.Max.Y-2, "p%.*f", prec, 100*q)
		}
		labelPixels++
	}

	// special case for first pixel
	drawf(d, 4, gdims.Max.Y-2, "p%.*f", 0, 100*pixel0Quantile)
}

func quantileAtPixel(n int) float64 {
	return 1 - (1-pixel0Quantile)*math.Pow(decayPerPixel, float64(n))
}

func valueAtPixel(hist *hdrhistogram.Histogram, n int) int64 {
	return hist.ValueAtQuantile(100 * quantileAtPixel(n))
}

func roundms(n int64) int64 {
	return int64(time.Duration(n) / time.Millisecond * time.Millisecond)
}

func digits(p float64) int {
	return -int(math.Floor(math.Log10(p)))
}

func drawf(d font.Drawer, x, y int, format string, args ...interface{}) {
	d.Dot = fixed.P(x, y)
	d.DrawString(fmt.Sprintf(format, args...))
}

func scale(v, from, to int64) int64 {
	if v > from {
		return to
	} else if v < 0 {
		return 0
	}
	return v * to / from
}

func pngError(w http.ResponseWriter, err error) {
	img := newImage(dims)
	d := drawer
	d.Dst = img
	drawf(d, 10, 25, "%s", err.Error())
	w.WriteHeader(500)
	png.Encode(w, img)
}

func newImage(dims image.Rectangle) *image.RGBA {
	img := image.NewRGBA(dims)
	draw.Draw(img, dims, image.Black, image.ZP, draw.Over)
	draw.Draw(img, dims.Inset(1), image.White, image.ZP, draw.Over)
	return img
}
<commit_msg>cmd/perfdash: show total number of events recorded<commit_after>package main

import (
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"math"
	"net/http"
	"strconv"
	"time"

	"github.com/codahale/hdrhistogram"

	"golang.org/x/image/font"
	"golang.org/x/image/font/inconsolata"
	"golang.org/x/image/math/fixed"
)

// TODO(kr): collect more fine-grain stats

const (
	// Use these to go to p99.99999
	// pixel0Quantile = 0.5
	// decayPerPixel  = 0.98
	pixel0Quantile = 0.4
	decayPerPixel  = 0.99
)

var (
	dims = image.Rect(0, 0, 800, 100)

	// In reverse order, newest to oldest,
	// so strongest to faintest color.
	// Don't forget, color.RGBA fields are
	// alpha premultiplied.
	lineColors = [...]color.Color{
		color.RGBA{R: 0xff, A: 0xff},
		color.RGBA{R: 0xcc, A: 0xcc},
		color.RGBA{R: 0x88, A: 0x88},
		color.RGBA{R: 0x55, A: 0x55},
		color.RGBA{R: 0x22, A: 0x22},
	}
	curLineColor = color.RGBA{A: 0x55}
	labelColor   = color.Black
	drawer       = font.Drawer{
		Src:  image.NewUniform(labelColor),
		Face: inconsolata.Regular8x16,
	}
)

func heatmap(w http.ResponseWriter, req *http.Request) {
	query := req.URL.Query()
	name := query.Get("name")
	id, err := strconv.Atoi(query.Get("id"))
	if err != nil {
		pngError(w, err)
		return
	}

	dv := getDebugVars(id)
	if dv == nil {
		pngError(w, errors.New("not found"))
		return
	}

	latency, ok := dv.Latency[name]
	if !ok {
		pngError(w, errors.New("not found"))
		return
	}

	img := newImage(dims)
	d := drawer
	d.Dst = img

	var (
		hists   []*hdrhistogram.Histogram
		trueMax time.Duration
		over    int
		total   int64
	)
	for _, b := range latency.Buckets {
		hists = append(hists, hdrhistogram.Import(&b.Histogram))
		over += b.Over
		total += int64(b.Over)
		if d := time.Duration(b.Max); d > trueMax {
			trueMax = d
		}
	}

	var max int64
	for _, hist := range hists {
		if v := hist.Max(); v > max {
			max = v
		}
		total += hist.TotalCount()
	}

	histMax := "no histogram data"
	if max > 0 {
		max += max / 10
		max = roundms(max)
		if max < int64(10*time.Millisecond) {
			max += int64(time.Millisecond)
		}
		complete := hists[:len(hists)-1]
		for i, hist := range complete {
			rindex := len(complete) - i - 1
			graph(img, hist, max, lineColors[rindex%len(lineColors)])
		}
		// special color for incomplete bucket
		graph(img, hists[len(hists)-1], max, curLineColor)
		label(img, labelColor)
		histMax = time.Duration(max).String()
	}
	drawf(d, 4, 20, "%s (%v max)", histMax, trueMax)
	drawf(d, 4, 38, "%s", name)
	drawf(d, 4, 54, "%d events (%d over)", total, over)
	png.Encode(w, img)
}

func graph(img *image.RGBA, hist *hdrhistogram.Histogram, ymax int64, color color.Color) {
	graph := img.SubImage(img.Bounds().Inset(2)).(*image.RGBA)
	gdims := graph.Bounds()
	d := drawer
	d.Dst = graph

	labelPixels := 50
	prevY := int(scale(valueAtPixel(hist, 0), ymax, int64(gdims.Dy())))
	for x := 1; x < gdims.Dx(); x++ {
		v := valueAtPixel(hist, x)
		y := int(scale(v, ymax, int64(gdims.Dy())))
		vLineSeg(graph, gdims.Min.X+x, gdims.Max.Y-y-1, gdims.Max.Y-prevY-1, color)
		labelPixels++
		prevY = y
	}
}

func vLineSeg(img *image.RGBA, x, y0, y1 int, color color.Color) {
	dims := image.Rect(x, y0, x+1, y1+1)
	draw.Draw(img, dims, image.NewUniform(color), image.ZP, draw.Over)
}

func label(img *image.RGBA, color color.Color) {
	graph := img.SubImage(img.Bounds().Inset(2)).(*image.RGBA)
	gdims := graph.Bounds()
	d := drawer
	d.Dst = graph

	labelDigits := 0
	labelPixels := 0
	for x := 0; x < gdims.Dx(); x++ {
		q := quantileAtPixel(x)
		if dig := digits(1 - q); labelPixels >= 50 && dig > labelDigits {
			labelPixels = 0
			labelDigits = dig
			for i := 0; i < 5; i++ {
				graph.Set(gdims.Min.X+x, gdims.Max.Y-i-1, color)
			}
			prec := labelDigits - 3
			if prec < 0 {
				prec = 0
			}
			drawf(d, gdims.Min.X+x+2, gdims.Max.Y-2, "p%.*f", prec, 100*q)
		}
		labelPixels++
	}

	// special case for first pixel
	drawf(d, 4, gdims.Max.Y-2, "p%.*f", 0, 100*pixel0Quantile)
}

func quantileAtPixel(n int) float64 {
	return 1 - (1-pixel0Quantile)*math.Pow(decayPerPixel, float64(n))
}

func valueAtPixel(hist *hdrhistogram.Histogram, n int) int64 {
	return hist.ValueAtQuantile(100 * quantileAtPixel(n))
}

func roundms(n int64) int64 {
	return int64(time.Duration(n) / time.Millisecond * time.Millisecond)
}

func digits(p float64) int {
	return -int(math.Floor(math.Log10(p)))
}

func drawf(d font.Drawer, x, y int, format string, args ...interface{}) {
	d.Dot = fixed.P(x, y)
	d.DrawString(fmt.Sprintf(format, args...))
}

func scale(v, from, to int64) int64 {
	if v > from {
		return to
	} else if v < 0 {
		return 0
	}
	return v * to / from
}

func pngError(w http.ResponseWriter, err error) {
	img := newImage(dims)
	d := drawer
	d.Dst = img
	drawf(d, 10, 25, "%s", err.Error())
	w.WriteHeader(500)
	png.Encode(w, img)
}

func newImage(dims image.Rectangle) *image.RGBA {
	img := image.NewRGBA(dims)
	draw.Draw(img, dims, image.Black, image.ZP, draw.Over)
	draw.Draw(img, dims.Inset(1), image.White, image.ZP, draw.Over)
	return img
}
<|endoftext|>


<commit_before>package geometry

import (
	"math"

	"github.com/gonum/matrix/mat64"
)

// Dot berechnet das Dot-Produkt aus a und b.
func Dot(a, b *mat64.Vector) float64 {
	return a.At(0, 0)*b.At(0, 0) + a.At(1, 0)*b.At(1, 0) + a.At(2, 0) + b.At(2, 0)
}

// Length berechnet die Länge eines Vectors
func Length(a *mat64.Vector) float64 {
	return math.Sqrt(a.At(0, 0)*a.At(0, 0) + a.At(1, 0)*a.At(1, 0) + a.At(2, 0)*a.At(2, 0))
}

// Normalize normiert einen Vektor auf die Länge 1
func Normalize(a *mat64.Vector) {
	length := Length(a)
	if length > 0 {
		a.ScaleVec(1/length, a)
	}
}
<commit_msg>Bug in Dot Product<commit_after>package geometry

import (
	"math"

	"github.com/gonum/matrix/mat64"
)

// Dot berechnet das Dot-Produkt aus a und b.
func Dot(a, b *mat64.Vector) float64 {
	return a.At(0, 0)*b.At(0, 0) + a.At(1, 0)*b.At(1, 0) + a.At(2, 0)*b.At(2, 0)
}

// Length berechnet die Länge eines Vectors
func Length(a *mat64.Vector) float64 {
	return math.Sqrt(a.At(0, 0)*a.At(0, 0) + a.At(1, 0)*a.At(1, 0) + a.At(2, 0)*a.At(2, 0))
}

// Normalize normiert einen Vektor auf die Länge 1
func Normalize(a *mat64.Vector) {
	length := Length(a)
	if length > 0 {
		a.ScaleVec(1/length, a)
	}
}
<|endoftext|>


<commit_before>// khan
// https://github.com/topfreegames/khan
//
// Licensed under the MIT license:
// http://www.opensource.org/licenses/mit-license
// Copyright © 2016 Top Free Games <backend@tfgco.com>

package util

// VERSION identifies Khan's current version
var VERSION = "4.1.3"
<commit_msg>Release v4.1.4<commit_after>// khan
// https://github.com/topfreegames/khan
//
// Licensed under the MIT license:
// http://www.opensource.org/licenses/mit-license
// Copyright © 2016 Top Free Games <backend@tfgco.com>

package util

// VERSION identifies Khan's current version
var VERSION = "4.1.4"
<|endoftext|>


<commit_before>package utils

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/spf13/cobra"
)

// Assumption describes the parameters that result in a Role
type Assumption struct {
	RoleName    string
	AccountID   string
	SessionName string
	UseMfa      bool
	MfaCode     string
	Lifetime    int64
	Policy      string
}

// AddAssumeFlags adds the flags for role assumption to a command object
func AddAssumeFlags(cmd *cobra.Command) {
	cmd.Flags().StringP("account", "a", "", "Account ID to assume role on (defaults to source account")
	cmd.Flags().StringP("session", "s", "", "Set session name for assumed role (defaults to origin user name)")
	cmd.Flags().Int64P("lifetime", "l", 3600, "Set lifetime of credentials in seconds (defaults to 3600 seconds / 1 hour, min 900, max 3600)")
	cmd.Flags().String("policy", "", "Set a IAM policy in JSON for the assumed credentials")
	cmd.Flags().BoolP("mfa", "m", false, "Use MFA when assuming role")
	cmd.Flags().String("mfacode", "", "Code to use for MFA")
}

// ParseFlags for assumption object
func (a *Assumption) ParseAssumeFlags(cmd *cobra.Command) error {
	flags := cmd.Flags()
	var err error
	a.AccountID, err = flags.GetString("account")
	if err != nil {
		return err
	}
	a.SessionName, err = flags.GetString("session")
	if err != nil {
		return err
	}
	a.Lifetime, err = flags.GetInt64("lifetime")
	if err != nil {
		return err
	}
	if a.Lifetime < 900 || a.Lifetime > 3600 {
		return fmt.Errorf("Lifetime must be between 900 and 3600: %d", a.Lifetime)
	}
	a.Policy, err = flags.GetString("policy")
	if err != nil {
		return err
	}
	a.UseMfa, err = flags.GetBool("mfa")
	if err != nil {
		return err
	}
	a.MfaCode, err = flags.GetString("mfacode")
	if err != nil {
		return err
	}
	return nil
}

// AssumeRole actions a role assumption object
func (a *Assumption) AssumeRole() (Role, error) {
	arn, err := API.RoleArn(a.RoleName, a.AccountID)
	if err != nil {
		return Role{}, err
	}
	if a.SessionName == "" {
		a.SessionName, err = API.SessionName()
		if err != nil {
			return Role{}, err
		}
	}

	params := &sts.AssumeRoleInput{
		RoleArn:         aws.String(arn),
		RoleSessionName: aws.String(a.SessionName),
		DurationSeconds: aws.Int64(a.Lifetime),
	}
	if a.Policy != "" {
		params.Policy = aws.String(a.Policy)
	}
	if err := configureMfa(a, params); err != nil {
		return Role{}, err
	}

	client := API.Client()
	resp, err := client.AssumeRole(params)
	if err != nil {
		return Role{}, err
	}

	creds := resp.Credentials
	newRole := Role{
		AccessKey:    *creds.AccessKeyId,
		SecretKey:    *creds.SecretAccessKey,
		SessionToken: *creds.SessionToken,
	}
	return newRole, nil
}
<commit_msg>rename assumeflags<commit_after>package utils

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/spf13/cobra"
)

// Assumption describes the parameters that result in a Role
type Assumption struct {
	RoleName    string
	AccountID   string
	SessionName string
	UseMfa      bool
	MfaCode     string
	Lifetime    int64
	Policy      string
}

// AddAssumeFlags adds the flags for role assumption to a command object
func AddAssumeFlags(cmd *cobra.Command) {
	cmd.Flags().StringP("account", "a", "", "Account ID to assume role on (defaults to source account")
	cmd.Flags().StringP("session", "s", "", "Set session name for assumed role (defaults to origin user name)")
	cmd.Flags().Int64P("lifetime", "l", 3600, "Set lifetime of credentials in seconds (defaults to 3600 seconds / 1 hour, min 900, max 3600)")
	cmd.Flags().String("policy", "", "Set a IAM policy in JSON for the assumed credentials")
	cmd.Flags().BoolP("mfa", "m", false, "Use MFA when assuming role")
	cmd.Flags().String("mfacode", "", "Code to use for MFA")
}

// ParseAssumeFlags for assumption object
func (a *Assumption) ParseAssumeFlags(cmd *cobra.Command) error {
	flags := cmd.Flags()
	var err error
	a.AccountID, err = flags.GetString("account")
	if err != nil {
		return err
	}
	a.SessionName, err = flags.GetString("session")
	if err != nil {
		return err
	}
	a.Lifetime, err = flags.GetInt64("lifetime")
	if err != nil {
		return err
	}
	if a.Lifetime < 900 || a.Lifetime > 3600 {
		return fmt.Errorf("Lifetime must be between 900 and 3600: %d", a.Lifetime)
	}
	a.Policy, err = flags.GetString("policy")
	if err != nil {
		return err
	}
	a.UseMfa, err = flags.GetBool("mfa")
	if err != nil {
		return err
	}
	a.MfaCode, err = flags.GetString("mfacode")
	if err != nil {
		return err
	}
	return nil
}

// AssumeRole actions a role assumption object
func (a *Assumption) AssumeRole() (Role, error) {
	arn, err := API.RoleArn(a.RoleName, a.AccountID)
	if err != nil {
		return Role{}, err
	}
	if a.SessionName == "" {
		a.SessionName, err = API.SessionName()
		if err != nil {
			return Role{}, err
		}
	}

	params := &sts.AssumeRoleInput{
		RoleArn:         aws.String(arn),
		RoleSessionName: aws.String(a.SessionName),
		DurationSeconds: aws.Int64(a.Lifetime),
	}
	if a.Policy != "" {
		params.Policy = aws.String(a.Policy)
	}
	if err := configureMfa(a, params); err != nil {
		return Role{}, err
	}

	client := API.Client()
	resp, err := client.AssumeRole(params)
	if err != nil {
		return Role{}, err
	}

	creds := resp.Credentials
	newRole := Role{
		AccessKey:    *creds.AccessKeyId,
		SecretKey:    *creds.SecretAccessKey,
		SessionToken: *creds.SessionToken,
	}
	return newRole, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus Authors
// 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.

// +build !nonetdev

package collector

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"

	"github.com/prometheus/common/log"
)

var (
	procNetDevFieldSep = regexp.MustCompile("[ :] *")
)

func getNetDevStats(ignore *regexp.Regexp) (map[string]map[string]string, error) {
	file, err := os.Open(procFilePath("net/dev"))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	return parseNetDevStats(file, ignore)
}

func parseNetDevStats(r io.Reader, ignore *regexp.Regexp) (map[string]map[string]string, error) {
	scanner := bufio.NewScanner(r)
	scanner.Scan() // skip first header
	scanner.Scan()
	parts := strings.Split(string(scanner.Text()), "|")
	if len(parts) != 3 { // interface + receive + transmit
		return nil, fmt.Errorf("invalid header line in net/dev: %s",
			scanner.Text())
	}

	header := strings.Fields(parts[1])
	netDev := map[string]map[string]string{}
	for scanner.Scan() {
		line := strings.TrimLeft(string(scanner.Text()), " ")
		parts := procNetDevFieldSep.Split(line, -1)
		if len(parts) != 2*len(header)+1 {
			return nil, fmt.Errorf("invalid line in net/dev: %s", scanner.Text())
		}

		dev := parts[0][:len(parts[0])]
		if ignore.MatchString(dev) {
			log.Debugf("Ignoring device: %s", dev)
			continue
		}
		netDev[dev] = map[string]string{}
		for i, v := range header {
			netDev[dev]["receive_"+v] = parts[i+1]
			netDev[dev]["transmit_"+v] = parts[i+1+len(header)]
		}
	}
	return netDev, nil
}
<commit_msg>Check for errors in netdev scanner<commit_after>// Copyright 2015 The Prometheus Authors
// 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.

// +build !nonetdev

package collector

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"

	"github.com/prometheus/common/log"
)

var (
	procNetDevFieldSep = regexp.MustCompile("[ :] *")
)

func getNetDevStats(ignore *regexp.Regexp) (map[string]map[string]string, error) {
	file, err := os.Open(procFilePath("net/dev"))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	return parseNetDevStats(file, ignore)
}

func parseNetDevStats(r io.Reader, ignore *regexp.Regexp) (map[string]map[string]string, error) {
	scanner := bufio.NewScanner(r)
	scanner.Scan() // skip first header
	scanner.Scan()
	parts := strings.Split(string(scanner.Text()), "|")
	if len(parts) != 3 { // interface + receive + transmit
		return nil, fmt.Errorf("invalid header line in net/dev: %s",
			scanner.Text())
	}

	header := strings.Fields(parts[1])
	netDev := map[string]map[string]string{}
	for scanner.Scan() {
		line := strings.TrimLeft(string(scanner.Text()), " ")
		parts := procNetDevFieldSep.Split(line, -1)
		if len(parts) != 2*len(header)+1 {
			return nil, fmt.Errorf("invalid line in net/dev: %s", scanner.Text())
		}

		dev := parts[0][:len(parts[0])]
		if ignore.MatchString(dev) {
			log.Debugf("Ignoring device: %s", dev)
			continue
		}
		netDev[dev] = map[string]string{}
		for i, v := range header {
			netDev[dev]["receive_"+v] = parts[i+1]
			netDev[dev]["transmit_"+v] = parts[i+1+len(header)]
		}
	}
	return netDev, scanner.Err()
}
<|endoftext|>


<commit_before>package command

import (
	"strings"
	"github.com/webdevops/go-shell"
	"github.com/webdevops/go-shell/commandbuilder"
	"fmt"
)

type PostgresCommonOptions struct {
	Hostname string `long:"hostname"`
	Username string `short:"u" long:"user"`
	Password string `short:"p" long:"password"`
	Docker   string `          long:"docker"`
	SSH      string `          long:"ssh"`

	connection commandbuilder.Connection
	dumpCompression string
}

func postgresQuote(value string) string {
	return "'" + strings.Replace(value, "'", "\\'", -1) + "'"
}

func postgresIdentifier(value string) string {
	return "\"" + strings.Replace(value, "\"", "\\\"", -1) + "\""
}

func  (conf *PostgresCommonOptions) Init() {
	if conf.SSH != "" {
		conf.connection.Hostname = conf.SSH
		fmt.Println(fmt.Sprintf(" - Using ssh connection \"%s\"", conf.SSH))
	}

	if conf.Docker != "" {
		conf.connection.Docker = conf.Docker
		conf.InitDockerSettings()
	}
}

func (conf *PostgresCommonOptions) PsqlCommandBuilder(args ...string) []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	cmd = append(cmd, "psql")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	if len(args) > 0 {
		cmd = append(cmd, args...)
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) PgDumpCommandBuilder(schema string) []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	cmd = append(cmd, "pg_dump")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	cmd = append(cmd, shell.Quote(schema))

	switch conf.dumpCompression {
	case "gzip":
		cmd = append(cmd, "| gzip")
	case "bzip2":
		cmd = append(cmd, "| bzip2")
	case "xz":
		cmd = append(cmd, "| xz --compress --stdout")
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) PgDumpAllCommandBuilder() []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	cmd = append(cmd, "pg_dumpall", "-c")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	switch conf.dumpCompression {
	case "gzip":
		cmd = append(cmd, "| gzip")
	case "bzip2":
		cmd = append(cmd, "| bzip2")
	case "xz":
		cmd = append(cmd, "| xz --compress --stdout")
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) PostgresRestoreCommandBuilder(args ...string) []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	switch conf.dumpCompression {
	case "gzip":
		cmd = append(cmd, "gzip -dc |")
	case "bzip2":
		cmd = append(cmd, "bzcat |")
	case "xz":
		cmd = append(cmd, "xzcat |")
	}

	cmd = append(cmd, "pg_dump")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	if len(args) > 0 {
		cmd = append(cmd, args...)
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) ExecStatement(statement string) string {
	cmd := shell.Cmd(conf.PsqlCommandBuilder("postgres", "-c", shell.Quote(statement))...)
	return cmd.Run().Stdout.String()
}

func  (conf *PostgresCommonOptions) InitDockerSettings() {
	containerName := conf.connection.Docker

	connectionClone := conf.connection.Clone()
	connectionClone.Docker = ""
	connectionClone.Type  = "auto"

	containerId := connectionClone.DockerGetContainerId(containerName)
	fmt.Println(fmt.Sprintf(" - Using docker container \"%s\"", containerId))

	containerEnv := GetDockerEnvList(connectionClone, containerId)

	// get user from env
	if val, ok := containerEnv["POSTGRES_USER"]; ok {
		if conf.Username == "" {
			conf.Username = val
		}
	}

	// get user from env
	if val, ok := containerEnv["POSTGRES_PASSWORD"]; ok {
		if conf.Username == "" {
			conf.Password = val
		}
	}
}
<commit_msg>Add port for postgres commands<commit_after>package command

import (
	"strings"
	"github.com/webdevops/go-shell"
	"github.com/webdevops/go-shell/commandbuilder"
	"fmt"
)

type PostgresCommonOptions struct {
	Hostname string `long:"hostname"`
	Port     string `short:"P" long:"port"`
	Username string `short:"u" long:"user"`
	Password string `short:"p" long:"password"`
	Docker   string `          long:"docker"`
	SSH      string `          long:"ssh"`

	connection commandbuilder.Connection
	dumpCompression string
}

func postgresQuote(value string) string {
	return "'" + strings.Replace(value, "'", "\\'", -1) + "'"
}

func postgresIdentifier(value string) string {
	return "\"" + strings.Replace(value, "\"", "\\\"", -1) + "\""
}

func  (conf *PostgresCommonOptions) Init() {
	if conf.SSH != "" {
		conf.connection.Hostname = conf.SSH
		fmt.Println(fmt.Sprintf(" - Using ssh connection \"%s\"", conf.SSH))
	}

	if conf.Docker != "" {
		conf.connection.Docker = conf.Docker
		conf.InitDockerSettings()
	}
}

func (conf *PostgresCommonOptions) PsqlCommandBuilder(args ...string) []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	cmd = append(cmd, "psql")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Port != "" {
		cmd = append(cmd, "-p", shell.Quote(conf.Port))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	if len(args) > 0 {
		cmd = append(cmd, args...)
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) PgDumpCommandBuilder(schema string) []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	cmd = append(cmd, "pg_dump")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Port != "" {
		cmd = append(cmd, "-p", shell.Quote(conf.Port))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	cmd = append(cmd, shell.Quote(schema))

	switch conf.dumpCompression {
	case "gzip":
		cmd = append(cmd, "| gzip")
	case "bzip2":
		cmd = append(cmd, "| bzip2")
	case "xz":
		cmd = append(cmd, "| xz --compress --stdout")
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) PgDumpAllCommandBuilder() []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	cmd = append(cmd, "pg_dumpall", "-c")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Port != "" {
		cmd = append(cmd, "-p", shell.Quote(conf.Port))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	switch conf.dumpCompression {
	case "gzip":
		cmd = append(cmd, "| gzip")
	case "bzip2":
		cmd = append(cmd, "| bzip2")
	case "xz":
		cmd = append(cmd, "| xz --compress --stdout")
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) PostgresRestoreCommandBuilder(args ...string) []interface{} {
	connection := conf.connection.Clone()
	cmd := []string{}

	switch conf.dumpCompression {
	case "gzip":
		cmd = append(cmd, "gzip -dc |")
	case "bzip2":
		cmd = append(cmd, "bzcat |")
	case "xz":
		cmd = append(cmd, "xzcat |")
	}

	cmd = append(cmd, "pg_dump")

	if conf.Hostname != "" {
		cmd = append(cmd, "-h", shell.Quote(conf.Hostname))
	}

	if conf.Port != "" {
		cmd = append(cmd, "-p", shell.Quote(conf.Port))
	}

	if conf.Username != "" {
		cmd = append(cmd, "-U", shell.Quote(conf.Username))
	}

	if conf.Password != "" {
		connection.Environment["PGPASSWORD"] = conf.Password
	}

	if len(args) > 0 {
		cmd = append(cmd, args...)
	}

	return connection.RawShellCommandBuilder(cmd...)
}

func (conf *PostgresCommonOptions) ExecStatement(statement string) string {
	cmd := shell.Cmd(conf.PsqlCommandBuilder("postgres", "-c", shell.Quote(statement))...)
	return cmd.Run().Stdout.String()
}

func  (conf *PostgresCommonOptions) InitDockerSettings() {
	containerName := conf.connection.Docker

	connectionClone := conf.connection.Clone()
	connectionClone.Docker = ""
	connectionClone.Type  = "auto"

	containerId := connectionClone.DockerGetContainerId(containerName)
	fmt.Println(fmt.Sprintf(" - Using docker container \"%s\"", containerId))

	containerEnv := GetDockerEnvList(connectionClone, containerId)

	// get user from env
	if val, ok := containerEnv["POSTGRES_USER"]; ok {
		if conf.Username == "" {
			conf.Username = val
		}
	}

	// get user from env
	if val, ok := containerEnv["POSTGRES_PASSWORD"]; ok {
		if conf.Username == "" {
			conf.Password = val
		}
	}
}
<|endoftext|>