<commit_before>package speed

import (
	"errors"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/performancecopilot/speed/bytebuffer"
)

// byte lengths of different components in an mmv file
const (
	HeaderLength         = 40
	TocLength            = 16
	MetricLength         = 104
	ValueLength          = 32
	InstanceLength       = 80
	InstanceDomainLength = 32
	StringBlockLength    = 256
)

// MaxMetricNameLength is the maximum length for a metric name
const MaxMetricNameLength = 63

// MaxDataValueSize is the maximum byte length for a stored metric value, unless it is a string
const MaxDataValueSize = 16

// EraseFileOnStop if set to true, will also delete the memory mapped file
var EraseFileOnStop = false

// Writer defines the interface of a MMV file writer's properties
type Writer interface {
	// a writer must contain a registry of metrics and instance domains
	Registry() Registry

	// writes an mmv file
	Start() error

	// stops writing and cleans up
	Stop() error

	// adds a metric to be written
	Register(Metric) error

	// adds an instance domain to be written
	RegisterIndom(InstanceDomain) error

	// adds metric and instance domain from a string
	RegisterString(string, interface{}, MetricSemantics, MetricType, MetricUnit) error

	// update a metric
	Update(Metric, interface{}) error
}

func mmvFileLocation(name string) (string, error) {
	if strings.ContainsRune(name, os.PathSeparator) {
		return "", errors.New("name cannot have path separator")
	}

	tdir, present := config["PCP_TMP_DIR"]
	var loc string
	if present {
		loc = path.Join(rootPath, tdir)
	} else {
		loc = os.TempDir()
	}

	return path.Join(loc, "mmv", name), nil
}

// PCPClusterIDBitLength is the bit length of the cluster id
// for a set of PCP metrics
const PCPClusterIDBitLength = 12

// MMVFlag represents an enumerated type to represent mmv flag values
type MMVFlag int

// values for MMVFlag
const (
	NoPrefixFlag MMVFlag = 1 << iota
	ProcessFlag
	SentinelFlag
)

//go:generate stringer -type=MMVFlag

// PCPWriter implements a writer that can write PCP compatible MMV files
type PCPWriter struct {
	sync.Mutex
	loc       string            // absolute location of the mmv file
	clusterID uint32            // cluster identifier for the writer
	flag      MMVFlag           // write flag
	r         *PCPRegistry      // current registry
	buffer    bytebuffer.Buffer // current Buffer
}

// NewPCPWriter initializes a new PCPWriter object
func NewPCPWriter(name string, flag MMVFlag) (*PCPWriter, error) {
	fileLocation, err := mmvFileLocation(name)
	if err != nil {
		return nil, err
	}

	return &PCPWriter{
		loc:       fileLocation,
		r:         NewPCPRegistry(),
		clusterID: getHash(name, PCPClusterIDBitLength),
		flag:      flag,
		buffer:    nil,
	}, nil
}

// Registry returns a writer's registry
func (w *PCPWriter) Registry() Registry {
	return w.r
}

func (w *PCPWriter) tocCount() int {
	ans := 2

	if w.Registry().InstanceCount() > 0 {
		ans += 2
	}

	if w.Registry().StringCount() > 0 {
		ans++
	}

	return ans
}

// Length returns the byte length of data in the mmv file written by the current writer
func (w *PCPWriter) Length() int {
	return HeaderLength +
		(w.tocCount() * TocLength) +
		(w.Registry().InstanceCount() * InstanceLength) +
		(w.Registry().InstanceDomainCount() * InstanceDomainLength) +
		(w.Registry().MetricCount() * (MetricLength + ValueLength)) +
		(w.Registry().StringCount() * StringBlockLength)
}

func (w *PCPWriter) initializeOffsets() {
	indomoffset := HeaderLength + TocLength*w.tocCount()
	instanceoffset := indomoffset + InstanceDomainLength*w.r.InstanceDomainCount()
	metricsoffset := instanceoffset + InstanceLength*w.r.InstanceCount()
	valuesoffset := metricsoffset + MetricLength*w.r.MetricCount()
	stringsoffset := valuesoffset + ValueLength*w.r.MetricCount()

	w.r.indomoffset = indomoffset
	w.r.instanceoffset = instanceoffset
	w.r.metricsoffset = metricsoffset
	w.r.valuesoffset = valuesoffset
	w.r.stringsoffset = stringsoffset

	for _, indom := range w.r.instanceDomains {
		indom.offset = indomoffset
		indom.instanceOffset = instanceoffset
		indomoffset += InstanceDomainLength

		for _, i := range indom.instances {
			i.offset = instanceoffset
			instanceoffset += InstanceLength
		}

		if indom.shortHelpText.val != "" {
			indom.shortHelpText.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if indom.longHelpText.val != "" {
			indom.longHelpText.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}

	for _, metric := range w.r.metrics {
		metric.desc.offset = metricsoffset
		metricsoffset += MetricLength
		metric.offset = valuesoffset
		valuesoffset += ValueLength

		if metric.desc.shortDescription.val != "" {
			metric.desc.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if metric.desc.longDescription.val != "" {
			metric.desc.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}
}

func (w *PCPWriter) writeHeaderBlock() (generation2offset int, generation int64) {
	// tag
	w.buffer.WriteString("MMV")
	w.buffer.SetPos(w.buffer.Pos() + 1) // extra null byte is needed and \0 isn't a valid escape character in go

	// version
	w.buffer.WriteUint32(1)

	// generation
	generation = time.Now().Unix()
	w.buffer.WriteInt64(generation)

	generation2offset = w.buffer.Pos()

	w.buffer.WriteInt64(0)

	// tocCount
	w.buffer.WriteInt32(int32(w.tocCount()))

	// flag mask
	w.buffer.WriteInt32(int32(w.flag))

	// process identifier
	w.buffer.WriteInt32(int32(os.Getpid()))

	// cluster identifier
	w.buffer.WriteUint32(w.clusterID)

	return
}

func (w *PCPWriter) writeSingleToc(pos, identifier, count, offset int) {
	w.buffer.SetPos(pos)
	w.buffer.WriteInt32(int32(identifier))
	w.buffer.WriteInt32(int32(count))
	w.buffer.WriteUint64(uint64(offset))
}

func (w *PCPWriter) writeTocBlock() {
	tocpos := HeaderLength

	// instance domains toc
	if w.Registry().InstanceDomainCount() > 0 {
		// 1 is the identifier for instance domains
		w.writeSingleToc(tocpos, 1, w.r.InstanceDomainCount(), w.r.indomoffset)
		tocpos += TocLength
	}

	// instances toc
	if w.Registry().InstanceCount() > 0 {
		// 2 is the identifier for instances
		w.writeSingleToc(tocpos, 2, w.r.InstanceCount(), w.r.instanceoffset)
		tocpos += TocLength
	}

	// metrics and values toc
	metricsoffset, valuesoffset := w.r.metricsoffset, w.r.valuesoffset
	if w.Registry().MetricCount() == 0 {
		metricsoffset, valuesoffset = 0, 0
	}

	// 3 is the identifier for metrics
	w.writeSingleToc(tocpos, 3, w.r.MetricCount(), metricsoffset)
	tocpos += TocLength

	// 4 is the identifier for values
	w.writeSingleToc(tocpos, 4, w.r.MetricCount(), valuesoffset)
	tocpos += TocLength

	// strings toc
	if w.Registry().StringCount() > 0 {
		// 5 is the identifier for strings
		w.writeSingleToc(tocpos, 5, w.r.StringCount(), w.r.stringsoffset)
	}
}

func (w *PCPWriter) writeInstanceAndInstanceDomainBlock() {
	for _, indom := range w.r.instanceDomains {
		w.buffer.SetPos(indom.offset)
		w.buffer.WriteUint32(indom.ID())
		w.buffer.WriteInt32(int32(indom.InstanceCount()))
		w.buffer.WriteInt64(int64(indom.instanceOffset))

		so, lo := indom.shortHelpText.offset, indom.longHelpText.offset
		w.buffer.WriteInt64(int64(so))
		w.buffer.WriteInt64(int64(lo))

		if so != 0 {
			w.buffer.SetPos(so)
			w.buffer.WriteString(indom.shortHelpText.val)
		}

		if lo != 0 {
			w.buffer.SetPos(lo)
			w.buffer.WriteString(indom.longHelpText.val)
		}

		for _, i := range indom.instances {
			w.buffer.SetPos(i.offset)
			w.buffer.WriteInt64(int64(indom.offset))
			w.buffer.WriteInt32(0)
			w.buffer.WriteUint32(i.id)
			w.buffer.WriteString(i.name)
		}
	}
}

func (w *PCPWriter) writeMetricDesc(desc *pcpMetricDesc) {
	pos := desc.offset
	w.buffer.SetPos(pos)

	w.buffer.WriteString(desc.name)
	w.buffer.Write([]byte{0})
	w.buffer.SetPos(pos + MaxMetricNameLength + 1)
	w.buffer.WriteUint32(desc.id)
	w.buffer.WriteInt32(int32(desc.t))
	w.buffer.WriteInt32(int32(desc.sem))
	w.buffer.WriteUint32(desc.u.PMAPI())
	if desc.indom != nil {
		w.buffer.WriteUint32(desc.indom.ID())
	} else {
		w.buffer.WriteInt32(-1)
	}
	w.buffer.WriteInt32(0)

	so, lo := desc.shortDescription.offset, desc.longDescription.offset
	w.buffer.WriteInt64(int64(so))
	w.buffer.WriteInt64(int64(lo))

	if so != 0 {
		w.buffer.SetPos(so)
		w.buffer.WriteString(desc.shortDescription.val)
	}

	if lo != 0 {
		w.buffer.SetPos(lo)
		w.buffer.WriteString(desc.longDescription.val)
	}
}

func (w *PCPWriter) writeMetricVal(m *PCPMetric) {
	pos := m.offset
	w.buffer.SetPos(pos)

	m.desc.t.WriteVal(m.val, w.buffer)

	w.buffer.SetPos(pos + MaxDataValueSize)
	w.buffer.WriteInt64(int64(m.desc.offset))
	if m.desc.indom != nil {
		w.buffer.WriteInt64(int64(m.desc.indom.(*PCPInstanceDomain).instanceOffset))
	} else {
		w.buffer.WriteInt64(0)
	}
}

func (w *PCPWriter) writeMetricsAndValuesBlock() {
	for _, metric := range w.r.metrics {
		w.writeMetricDesc(metric.desc)
		w.writeMetricVal(metric)
	}
}

// fillData will fill the Buffer with the mmv file
// data as long as something doesn't go wrong
func (w *PCPWriter) fillData() error {
	generation2offset, generation := w.writeHeaderBlock()
	w.writeTocBlock()
	w.writeInstanceAndInstanceDomainBlock()
	w.writeMetricsAndValuesBlock()

	w.buffer.SetPos(generation2offset)
	w.buffer.WriteUint64(uint64(generation))

	return nil
}

// Start dumps existing registry data
func (w *PCPWriter) Start() error {
	w.Lock()
	defer w.Unlock()

	l := w.Length()

	w.initializeOffsets()

	buffer, err := bytebuffer.NewMemoryMappedBuffer(w.loc, l)
	if err != nil {
		return err
	}
	w.buffer = buffer

	w.fillData()

	w.r.mapped = true

	return nil
}

// Stop removes existing mapping and cleans up
func (w *PCPWriter) Stop() error {
	w.Lock()
	defer w.Unlock()

	w.r.mapped = false

	err := w.buffer.(*bytebuffer.MemoryMappedBuffer).Unmap(EraseFileOnStop)
	w.buffer = nil
	if err != nil {
		return err
	}

	return nil
}

// Register is simply a shorthand for Registry().AddMetric
func (w *PCPWriter) Register(m Metric) error { return w.Registry().AddMetric(m) }

// RegisterIndom is simply a shorthand for Registry().AddInstanceDomain
func (w *PCPWriter) RegisterIndom(indom InstanceDomain) error {
	return w.Registry().AddInstanceDomain(indom)
}

// RegisterString is simply a shorthand for Registry().AddMetricByString
func (w *PCPWriter) RegisterString(str string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) error {
	_, err := w.Registry().AddMetricByString(str, val, s, t, u)
	return err
}

// Update is simply a shorthand for Registry().UpdateMetric
func (w *PCPWriter) Update(m Metric, val interface{}) error { return w.Registry().UpdateMetric(m, val) }
<commit_msg>writer: add Length to the Writer API<commit_after>package speed

import (
	"errors"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/performancecopilot/speed/bytebuffer"
)

// byte lengths of different components in an mmv file
const (
	HeaderLength         = 40
	TocLength            = 16
	MetricLength         = 104
	ValueLength          = 32
	InstanceLength       = 80
	InstanceDomainLength = 32
	StringBlockLength    = 256
)

// MaxMetricNameLength is the maximum length for a metric name
const MaxMetricNameLength = 63

// MaxDataValueSize is the maximum byte length for a stored metric value, unless it is a string
const MaxDataValueSize = 16

// EraseFileOnStop if set to true, will also delete the memory mapped file
var EraseFileOnStop = false

// Writer defines the interface of a MMV file writer's properties
type Writer interface {
	// a writer must contain a registry of metrics and instance domains
	Registry() Registry

	// writes an mmv file
	Start() error

	// stops writing and cleans up
	Stop() error

	// returns the number of bytes that will be written by the current writer
	Length() int

	// adds a metric to be written
	Register(Metric) error

	// adds an instance domain to be written
	RegisterIndom(InstanceDomain) error

	// adds metric and instance domain from a string
	RegisterString(string, interface{}, MetricSemantics, MetricType, MetricUnit) error

	// update a metric
	Update(Metric, interface{}) error
}

func mmvFileLocation(name string) (string, error) {
	if strings.ContainsRune(name, os.PathSeparator) {
		return "", errors.New("name cannot have path separator")
	}

	tdir, present := config["PCP_TMP_DIR"]
	var loc string
	if present {
		loc = path.Join(rootPath, tdir)
	} else {
		loc = os.TempDir()
	}

	return path.Join(loc, "mmv", name), nil
}

// PCPClusterIDBitLength is the bit length of the cluster id
// for a set of PCP metrics
const PCPClusterIDBitLength = 12

// MMVFlag represents an enumerated type to represent mmv flag values
type MMVFlag int

// values for MMVFlag
const (
	NoPrefixFlag MMVFlag = 1 << iota
	ProcessFlag
	SentinelFlag
)

//go:generate stringer -type=MMVFlag

// PCPWriter implements a writer that can write PCP compatible MMV files
type PCPWriter struct {
	sync.Mutex
	loc       string            // absolute location of the mmv file
	clusterID uint32            // cluster identifier for the writer
	flag      MMVFlag           // write flag
	r         *PCPRegistry      // current registry
	buffer    bytebuffer.Buffer // current Buffer
}

// NewPCPWriter initializes a new PCPWriter object
func NewPCPWriter(name string, flag MMVFlag) (*PCPWriter, error) {
	fileLocation, err := mmvFileLocation(name)
	if err != nil {
		return nil, err
	}

	return &PCPWriter{
		loc:       fileLocation,
		r:         NewPCPRegistry(),
		clusterID: getHash(name, PCPClusterIDBitLength),
		flag:      flag,
		buffer:    nil,
	}, nil
}

// Registry returns a writer's registry
func (w *PCPWriter) Registry() Registry {
	return w.r
}

func (w *PCPWriter) tocCount() int {
	ans := 2

	if w.Registry().InstanceCount() > 0 {
		ans += 2
	}

	if w.Registry().StringCount() > 0 {
		ans++
	}

	return ans
}

// Length returns the byte length of data in the mmv file written by the current writer
func (w *PCPWriter) Length() int {
	return HeaderLength +
		(w.tocCount() * TocLength) +
		(w.Registry().InstanceCount() * InstanceLength) +
		(w.Registry().InstanceDomainCount() * InstanceDomainLength) +
		(w.Registry().MetricCount() * (MetricLength + ValueLength)) +
		(w.Registry().StringCount() * StringBlockLength)
}

func (w *PCPWriter) initializeOffsets() {
	indomoffset := HeaderLength + TocLength*w.tocCount()
	instanceoffset := indomoffset + InstanceDomainLength*w.r.InstanceDomainCount()
	metricsoffset := instanceoffset + InstanceLength*w.r.InstanceCount()
	valuesoffset := metricsoffset + MetricLength*w.r.MetricCount()
	stringsoffset := valuesoffset + ValueLength*w.r.MetricCount()

	w.r.indomoffset = indomoffset
	w.r.instanceoffset = instanceoffset
	w.r.metricsoffset = metricsoffset
	w.r.valuesoffset = valuesoffset
	w.r.stringsoffset = stringsoffset

	for _, indom := range w.r.instanceDomains {
		indom.offset = indomoffset
		indom.instanceOffset = instanceoffset
		indomoffset += InstanceDomainLength

		for _, i := range indom.instances {
			i.offset = instanceoffset
			instanceoffset += InstanceLength
		}

		if indom.shortDescription.val != "" {
			indom.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if indom.longDescription.val != "" {
			indom.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}

	for _, metric := range w.r.metrics {
		metric.desc.offset = metricsoffset
		metricsoffset += MetricLength
		metric.offset = valuesoffset
		valuesoffset += ValueLength

		if metric.desc.shortDescription.val != "" {
			metric.desc.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if metric.desc.longDescription.val != "" {
			metric.desc.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}
}

func (w *PCPWriter) writeHeaderBlock() (generation2offset int, generation int64) {
	// tag
	w.buffer.WriteString("MMV")
	w.buffer.SetPos(w.buffer.Pos() + 1) // extra null byte is needed and \0 isn't a valid escape character in go

	// version
	w.buffer.WriteUint32(1)

	// generation
	generation = time.Now().Unix()
	w.buffer.WriteInt64(generation)

	generation2offset = w.buffer.Pos()

	w.buffer.WriteInt64(0)

	// tocCount
	w.buffer.WriteInt32(int32(w.tocCount()))

	// flag mask
	w.buffer.WriteInt32(int32(w.flag))

	// process identifier
	w.buffer.WriteInt32(int32(os.Getpid()))

	// cluster identifier
	w.buffer.WriteUint32(w.clusterID)

	return
}

func (w *PCPWriter) writeSingleToc(pos, identifier, count, offset int) {
	w.buffer.SetPos(pos)
	w.buffer.WriteInt32(int32(identifier))
	w.buffer.WriteInt32(int32(count))
	w.buffer.WriteUint64(uint64(offset))
}

func (w *PCPWriter) writeTocBlock() {
	tocpos := HeaderLength

	// instance domains toc
	if w.Registry().InstanceDomainCount() > 0 {
		// 1 is the identifier for instance domains
		w.writeSingleToc(tocpos, 1, w.r.InstanceDomainCount(), w.r.indomoffset)
		tocpos += TocLength
	}

	// instances toc
	if w.Registry().InstanceCount() > 0 {
		// 2 is the identifier for instances
		w.writeSingleToc(tocpos, 2, w.r.InstanceCount(), w.r.instanceoffset)
		tocpos += TocLength
	}

	// metrics and values toc
	metricsoffset, valuesoffset := w.r.metricsoffset, w.r.valuesoffset
	if w.Registry().MetricCount() == 0 {
		metricsoffset, valuesoffset = 0, 0
	}

	// 3 is the identifier for metrics
	w.writeSingleToc(tocpos, 3, w.r.MetricCount(), metricsoffset)
	tocpos += TocLength

	// 4 is the identifier for values
	w.writeSingleToc(tocpos, 4, w.r.MetricCount(), valuesoffset)
	tocpos += TocLength

	// strings toc
	if w.Registry().StringCount() > 0 {
		// 5 is the identifier for strings
		w.writeSingleToc(tocpos, 5, w.r.StringCount(), w.r.stringsoffset)
	}
}

func (w *PCPWriter) writeInstanceAndInstanceDomainBlock() {
	for _, indom := range w.r.instanceDomains {
		w.buffer.SetPos(indom.offset)
		w.buffer.WriteUint32(indom.ID())
		w.buffer.WriteInt32(int32(indom.InstanceCount()))
		w.buffer.WriteInt64(int64(indom.instanceOffset))

		so, lo := indom.shortDescription.offset, indom.longDescription.offset
		w.buffer.WriteInt64(int64(so))
		w.buffer.WriteInt64(int64(lo))

		if so != 0 {
			w.buffer.SetPos(so)
			w.buffer.WriteString(indom.shortDescription.val)
		}

		if lo != 0 {
			w.buffer.SetPos(lo)
			w.buffer.WriteString(indom.longDescription.val)
		}

		for _, i := range indom.instances {
			w.buffer.SetPos(i.offset)
			w.buffer.WriteInt64(int64(indom.offset))
			w.buffer.WriteInt32(0)
			w.buffer.WriteUint32(i.id)
			w.buffer.WriteString(i.name)
		}
	}
}

func (w *PCPWriter) writeMetricDesc(desc *pcpMetricDesc) {
	pos := desc.offset
	w.buffer.SetPos(pos)

	w.buffer.WriteString(desc.name)
	w.buffer.Write([]byte{0})
	w.buffer.SetPos(pos + MaxMetricNameLength + 1)
	w.buffer.WriteUint32(desc.id)
	w.buffer.WriteInt32(int32(desc.t))
	w.buffer.WriteInt32(int32(desc.sem))
	w.buffer.WriteUint32(desc.u.PMAPI())
	if desc.indom != nil {
		w.buffer.WriteUint32(desc.indom.ID())
	} else {
		w.buffer.WriteInt32(-1)
	}
	w.buffer.WriteInt32(0)

	so, lo := desc.shortDescription.offset, desc.longDescription.offset
	w.buffer.WriteInt64(int64(so))
	w.buffer.WriteInt64(int64(lo))

	if so != 0 {
		w.buffer.SetPos(so)
		w.buffer.WriteString(desc.shortDescription.val)
	}

	if lo != 0 {
		w.buffer.SetPos(lo)
		w.buffer.WriteString(desc.longDescription.val)
	}
}

func (w *PCPWriter) writeMetricVal(m *PCPMetric) {
	pos := m.offset
	w.buffer.SetPos(pos)

	m.desc.t.WriteVal(m.val, w.buffer)

	w.buffer.SetPos(pos + MaxDataValueSize)
	w.buffer.WriteInt64(int64(m.desc.offset))
	if m.desc.indom != nil {
		w.buffer.WriteInt64(int64(m.desc.indom.(*PCPInstanceDomain).instanceOffset))
	} else {
		w.buffer.WriteInt64(0)
	}
}

func (w *PCPWriter) writeMetricsAndValuesBlock() {
	for _, metric := range w.r.metrics {
		w.writeMetricDesc(metric.desc)
		w.writeMetricVal(metric)
	}
}

// fillData will fill the Buffer with the mmv file
// data as long as something doesn't go wrong
func (w *PCPWriter) fillData() error {
	generation2offset, generation := w.writeHeaderBlock()
	w.writeTocBlock()
	w.writeInstanceAndInstanceDomainBlock()
	w.writeMetricsAndValuesBlock()

	w.buffer.SetPos(generation2offset)
	w.buffer.WriteUint64(uint64(generation))

	return nil
}

// Start dumps existing registry data
func (w *PCPWriter) Start() error {
	w.Lock()
	defer w.Unlock()

	l := w.Length()

	w.initializeOffsets()

	buffer, err := bytebuffer.NewMemoryMappedBuffer(w.loc, l)
	if err != nil {
		return err
	}
	w.buffer = buffer

	w.fillData()

	w.r.mapped = true

	return nil
}

// Stop removes existing mapping and cleans up
func (w *PCPWriter) Stop() error {
	w.Lock()
	defer w.Unlock()

	w.r.mapped = false

	err := w.buffer.(*bytebuffer.MemoryMappedBuffer).Unmap(EraseFileOnStop)
	w.buffer = nil
	if err != nil {
		return err
	}

	return nil
}

// Register is simply a shorthand for Registry().AddMetric
func (w *PCPWriter) Register(m Metric) error { return w.Registry().AddMetric(m) }

// RegisterIndom is simply a shorthand for Registry().AddInstanceDomain
func (w *PCPWriter) RegisterIndom(indom InstanceDomain) error {
	return w.Registry().AddInstanceDomain(indom)
}

// RegisterString is simply a shorthand for Registry().AddMetricByString
func (w *PCPWriter) RegisterString(str string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) error {
	_, err := w.Registry().AddMetricByString(str, val, s, t, u)
	return err
}

// Update is simply a shorthand for Registry().UpdateMetric
func (w *PCPWriter) Update(m Metric, val interface{}) error { return w.Registry().UpdateMetric(m, val) }
<|endoftext|>


<commit_before>// Copyright 2018 The ACH Authors
// Use of this source code is governed by an Apache License
// license that can be found in the LICENSE file.

package ach

import (
	"bufio"
	"io"
	"strings"
)

// A Writer writes an ach.file to a NACHA encoded file.
//
// As returned by NewWriter, a Writer writes ach.file structs into
// NACHA formatted files.
//
type Writer struct {
	w       *bufio.Writer
	lineNum int //current line being written
}

// NewWriter returns a new Writer that writes to w.
func NewWriter(w io.Writer) *Writer {
	return &Writer{
		w: bufio.NewWriter(w),
	}
}

// Writer writes a single ach.file record to w
func (w *Writer) Write(file *File) error {
	if err := file.Validate(); err != nil {
		return err
	}

	w.lineNum = 0
	// Iterate over all records in the file
	w.w.WriteString(file.Header.String() + "\n")
	w.lineNum++

	for _, batch := range file.Batches {
		w.w.WriteString(batch.GetHeader().String() + "\n")
		w.lineNum++
		for _, entry := range batch.GetEntries() {
			w.w.WriteString(entry.String() + "\n")
			w.lineNum++
			for _, addenda := range entry.Addendum {
				w.w.WriteString(addenda.String() + "\n")
				w.lineNum++
			}
		}
		w.w.WriteString(batch.GetControl().String() + "\n")
		w.lineNum++
	}
	w.w.WriteString(file.Control.String() + "\n")
	w.lineNum++

	// pad the final block
	for i := 0; i < (10-(w.lineNum%10)) && w.lineNum%10 != 0; i++ {
		w.w.WriteString(strings.Repeat("9", 94) + "\n")
	}

	return w.w.Flush()
}

// Flush writes any buffered data to the underlying io.Writer.
// To check if an error occurred during the Flush, call Error.
func (w *Writer) Flush() {
	w.w.Flush()
}
<commit_msg>Add back in error trapping on writes<commit_after>// Copyright 2018 The ACH Authors
// Use of this source code is governed by an Apache License
// license that can be found in the LICENSE file.

package ach

import (
	"bufio"
	"io"
	"strings"
)

// A Writer writes an ach.file to a NACHA encoded file.
//
// As returned by NewWriter, a Writer writes ach.file structs into
// NACHA formatted files.
//
type Writer struct {
	w       *bufio.Writer
	lineNum int //current line being written
}

// NewWriter returns a new Writer that writes to w.
func NewWriter(w io.Writer) *Writer {
	return &Writer{
		w: bufio.NewWriter(w),
	}
}

// Writer writes a single ach.file record to w
func (w *Writer) Write(file *File) error {
	if err := file.Validate(); err != nil {
		return err
	}

	w.lineNum = 0
	// Iterate over all records in the file
	if _, err := w.w.WriteString(file.Header.String() + "\n"); err != nil {
		return err
	}
	w.lineNum++

	for _, batch := range file.Batches {
		if _, err := w.w.WriteString(batch.GetHeader().String() + "\n"); err != nil {
			return err
		}
		w.lineNum++
		for _, entry := range batch.GetEntries() {
			if _, err := w.w.WriteString(entry.String() + "\n"); err != nil {
				return err
			}
			w.lineNum++
			for _, addenda := range entry.Addendum {
				if _, err := w.w.WriteString(addenda.String() + "\n"); err != nil {
					return err
				}
				w.lineNum++
			}
		}
		if _, err := w.w.WriteString(batch.GetControl().String() + "\n"); err != nil {
			return err
		}
		w.lineNum++
	}
	if _, err := w.w.WriteString(file.Control.String() + "\n"); err != nil {
		return err
	}
	w.lineNum++

	// pad the final block
	for i := 0; i < (10-(w.lineNum%10)) && w.lineNum%10 != 0; i++ {
		if _, err := w.w.WriteString(strings.Repeat("9", 94) + "\n"); err != nil {
			return err
		}
	}

	return w.w.Flush()
}

// Flush writes any buffered data to the underlying io.Writer.
// To check if an error occurred during the Flush, call Error.
func (w *Writer) Flush() {
	w.w.Flush()
}
<|endoftext|>


<commit_before>package speed

import (
	"errors"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/performancecopilot/speed/bytebuffer"
)

// byte lengths of different components in an mmv file
const (
	HeaderLength         = 40
	TocLength            = 16
	MetricLength         = 104
	ValueLength          = 32
	InstanceLength       = 80
	InstanceDomainLength = 32
	StringBlockLength    = 256
)

// MaxMetricNameLength is the maximum length for a metric name
const MaxMetricNameLength = 63

// MaxDataValueSize is the maximum byte length for a stored metric value, unless it is a string
const MaxDataValueSize = 16

// EraseFileOnStop if set to true, will also delete the memory mapped file
var EraseFileOnStop = false

// Writer defines the interface of a MMV file writer's properties
type Writer interface {
	// a writer must contain a registry of metrics and instance domains
	Registry() Registry

	// writes an mmv file
	Start() error

	// stops writing and cleans up
	Stop() error

	// returns the number of bytes that will be written by the current writer
	Length() int

	// adds a metric to be written
	Register(Metric) error

	// adds an instance domain to be written
	RegisterIndom(InstanceDomain) error

	// adds metric and instance domain from a string
	RegisterString(string, interface{}, MetricSemantics, MetricType, MetricUnit) error

	// update a metric
	Update(Metric, interface{}) error
}

func mmvFileLocation(name string) (string, error) {
	if strings.ContainsRune(name, os.PathSeparator) {
		return "", errors.New("name cannot have path separator")
	}

	tdir, present := config["PCP_TMP_DIR"]
	var loc string
	if present {
		loc = path.Join(rootPath, tdir)
	} else {
		loc = os.TempDir()
	}

	return path.Join(loc, "mmv", name), nil
}

// PCPClusterIDBitLength is the bit length of the cluster id
// for a set of PCP metrics
const PCPClusterIDBitLength = 12

// MMVFlag represents an enumerated type to represent mmv flag values
type MMVFlag int

// values for MMVFlag
const (
	NoPrefixFlag MMVFlag = 1 << iota
	ProcessFlag
	SentinelFlag
)

//go:generate stringer -type=MMVFlag

// PCPWriter implements a writer that can write PCP compatible MMV files
type PCPWriter struct {
	sync.Mutex
	loc       string            // absolute location of the mmv file
	clusterID uint32            // cluster identifier for the writer
	flag      MMVFlag           // write flag
	r         *PCPRegistry      // current registry
	buffer    bytebuffer.Buffer // current Buffer
}

// NewPCPWriter initializes a new PCPWriter object
func NewPCPWriter(name string, flag MMVFlag) (*PCPWriter, error) {
	fileLocation, err := mmvFileLocation(name)
	if err != nil {
		return nil, err
	}

	return &PCPWriter{
		loc:       fileLocation,
		r:         NewPCPRegistry(),
		clusterID: getHash(name, PCPClusterIDBitLength),
		flag:      flag,
		buffer:    nil,
	}, nil
}

// Registry returns a writer's registry
func (w *PCPWriter) Registry() Registry {
	return w.r
}

func (w *PCPWriter) tocCount() int {
	ans := 2

	if w.Registry().InstanceCount() > 0 {
		ans += 2
	}

	if w.Registry().StringCount() > 0 {
		ans++
	}

	return ans
}

// Length returns the byte length of data in the mmv file written by the current writer
func (w *PCPWriter) Length() int {
	return HeaderLength +
		(w.tocCount() * TocLength) +
		(w.Registry().InstanceCount() * InstanceLength) +
		(w.Registry().InstanceDomainCount() * InstanceDomainLength) +
		(w.Registry().MetricCount() * (MetricLength + ValueLength)) +
		(w.Registry().StringCount() * StringBlockLength)
}

func (w *PCPWriter) initializeOffsets() {
	indomoffset := HeaderLength + TocLength*w.tocCount()
	instanceoffset := indomoffset + InstanceDomainLength*w.r.InstanceDomainCount()
	metricsoffset := instanceoffset + InstanceLength*w.r.InstanceCount()
	valuesoffset := metricsoffset + MetricLength*w.r.MetricCount()
	stringsoffset := valuesoffset + ValueLength*w.r.MetricCount()

	w.r.indomoffset = indomoffset
	w.r.instanceoffset = instanceoffset
	w.r.metricsoffset = metricsoffset
	w.r.valuesoffset = valuesoffset
	w.r.stringsoffset = stringsoffset

	for _, indom := range w.r.instanceDomains {
		indom.offset = indomoffset
		indom.instanceOffset = instanceoffset
		indomoffset += InstanceDomainLength

		for _, i := range indom.instances {
			i.offset = instanceoffset
			instanceoffset += InstanceLength
		}

		if indom.shortDescription.val != "" {
			indom.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if indom.longDescription.val != "" {
			indom.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}

	for _, metric := range w.r.metrics {
		metric.desc.offset = metricsoffset
		metricsoffset += MetricLength
		metric.offset = valuesoffset
		valuesoffset += ValueLength

		if metric.desc.shortDescription.val != "" {
			metric.desc.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if metric.desc.longDescription.val != "" {
			metric.desc.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}
}

func (w *PCPWriter) writeHeaderBlock() (generation2offset int, generation int64) {
	// tag
	w.buffer.WriteString("MMV")
	w.buffer.SetPos(w.buffer.Pos() + 1) // extra null byte is needed and \0 isn't a valid escape character in go

	// version
	w.buffer.WriteUint32(1)

	// generation
	generation = time.Now().Unix()
	w.buffer.WriteInt64(generation)

	generation2offset = w.buffer.Pos()

	w.buffer.WriteInt64(0)

	// tocCount
	w.buffer.WriteInt32(int32(w.tocCount()))

	// flag mask
	w.buffer.WriteInt32(int32(w.flag))

	// process identifier
	w.buffer.WriteInt32(int32(os.Getpid()))

	// cluster identifier
	w.buffer.WriteUint32(w.clusterID)

	return
}

func (w *PCPWriter) writeSingleToc(pos, identifier, count, offset int) {
	w.buffer.SetPos(pos)
	w.buffer.WriteInt32(int32(identifier))
	w.buffer.WriteInt32(int32(count))
	w.buffer.WriteUint64(uint64(offset))
}

func (w *PCPWriter) writeTocBlock() {
	tocpos := HeaderLength

	// instance domains toc
	if w.Registry().InstanceDomainCount() > 0 {
		// 1 is the identifier for instance domains
		w.writeSingleToc(tocpos, 1, w.r.InstanceDomainCount(), w.r.indomoffset)
		tocpos += TocLength
	}

	// instances toc
	if w.Registry().InstanceCount() > 0 {
		// 2 is the identifier for instances
		w.writeSingleToc(tocpos, 2, w.r.InstanceCount(), w.r.instanceoffset)
		tocpos += TocLength
	}

	// metrics and values toc
	metricsoffset, valuesoffset := w.r.metricsoffset, w.r.valuesoffset
	if w.Registry().MetricCount() == 0 {
		metricsoffset, valuesoffset = 0, 0
	}

	// 3 is the identifier for metrics
	w.writeSingleToc(tocpos, 3, w.r.MetricCount(), metricsoffset)
	tocpos += TocLength

	// 4 is the identifier for values
	w.writeSingleToc(tocpos, 4, w.r.MetricCount(), valuesoffset)
	tocpos += TocLength

	// strings toc
	if w.Registry().StringCount() > 0 {
		// 5 is the identifier for strings
		w.writeSingleToc(tocpos, 5, w.r.StringCount(), w.r.stringsoffset)
	}
}

func (w *PCPWriter) writeInstanceAndInstanceDomainBlock() {
	for _, indom := range w.r.instanceDomains {
		w.buffer.SetPos(indom.offset)
		w.buffer.WriteUint32(indom.ID())
		w.buffer.WriteInt32(int32(indom.InstanceCount()))
		w.buffer.WriteInt64(int64(indom.instanceOffset))

		so, lo := indom.shortDescription.offset, indom.longDescription.offset
		w.buffer.WriteInt64(int64(so))
		w.buffer.WriteInt64(int64(lo))

		if so != 0 {
			w.buffer.SetPos(so)
			w.buffer.WriteString(indom.shortDescription.val)
		}

		if lo != 0 {
			w.buffer.SetPos(lo)
			w.buffer.WriteString(indom.longDescription.val)
		}

		for _, i := range indom.instances {
			w.buffer.SetPos(i.offset)
			w.buffer.WriteInt64(int64(indom.offset))
			w.buffer.WriteInt32(0)
			w.buffer.WriteUint32(i.id)
			w.buffer.WriteString(i.name)
		}
	}
}

func (w *PCPWriter) writeMetricDesc(desc *pcpMetricDesc) {
	pos := desc.offset
	w.buffer.SetPos(pos)

	w.buffer.WriteString(desc.name)
	w.buffer.Write([]byte{0})
	w.buffer.SetPos(pos + MaxMetricNameLength + 1)
	w.buffer.WriteUint32(desc.id)
	w.buffer.WriteInt32(int32(desc.t))
	w.buffer.WriteInt32(int32(desc.sem))
	w.buffer.WriteUint32(desc.u.PMAPI())
	if desc.indom != nil {
		w.buffer.WriteUint32(desc.indom.ID())
	} else {
		w.buffer.WriteInt32(-1)
	}
	w.buffer.WriteInt32(0)

	so, lo := desc.shortDescription.offset, desc.longDescription.offset
	w.buffer.WriteInt64(int64(so))
	w.buffer.WriteInt64(int64(lo))

	if so != 0 {
		w.buffer.SetPos(so)
		w.buffer.WriteString(desc.shortDescription.val)
	}

	if lo != 0 {
		w.buffer.SetPos(lo)
		w.buffer.WriteString(desc.longDescription.val)
	}
}

func (w *PCPWriter) writeMetricVal(m *PCPMetric) {
	pos := m.offset
	w.buffer.SetPos(pos)

	m.desc.t.WriteVal(m.val, w.buffer)

	w.buffer.SetPos(pos + MaxDataValueSize)
	w.buffer.WriteInt64(int64(m.desc.offset))
	if m.desc.indom != nil {
		w.buffer.WriteInt64(int64(m.desc.indom.(*PCPInstanceDomain).instanceOffset))
	} else {
		w.buffer.WriteInt64(0)
	}
}

func (w *PCPWriter) writeMetricsAndValuesBlock() {
	for _, metric := range w.r.metrics {
		w.writeMetricDesc(metric.desc)
		w.writeMetricVal(metric)
	}
}

// fillData will fill the Buffer with the mmv file
// data as long as something doesn't go wrong
func (w *PCPWriter) fillData() error {
	generation2offset, generation := w.writeHeaderBlock()
	w.writeTocBlock()
	w.writeInstanceAndInstanceDomainBlock()
	w.writeMetricsAndValuesBlock()

	w.buffer.SetPos(generation2offset)
	w.buffer.WriteUint64(uint64(generation))

	return nil
}

// Start dumps existing registry data
func (w *PCPWriter) Start() error {
	w.Lock()
	defer w.Unlock()

	l := w.Length()

	w.initializeOffsets()

	buffer, err := bytebuffer.NewMemoryMappedBuffer(w.loc, l)
	if err != nil {
		return err
	}
	w.buffer = buffer

	w.fillData()

	w.r.mapped = true

	return nil
}

// Stop removes existing mapping and cleans up
func (w *PCPWriter) Stop() error {
	w.Lock()
	defer w.Unlock()

	w.r.mapped = false

	err := w.buffer.(*bytebuffer.MemoryMappedBuffer).Unmap(EraseFileOnStop)
	w.buffer = nil
	if err != nil {
		return err
	}

	return nil
}

// Register is simply a shorthand for Registry().AddMetric
func (w *PCPWriter) Register(m Metric) error { return w.Registry().AddMetric(m) }

// RegisterIndom is simply a shorthand for Registry().AddInstanceDomain
func (w *PCPWriter) RegisterIndom(indom InstanceDomain) error {
	return w.Registry().AddInstanceDomain(indom)
}

// RegisterString is simply a shorthand for Registry().AddMetricByString
func (w *PCPWriter) RegisterString(str string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) error {
	_, err := w.Registry().AddMetricByString(str, val, s, t, u)
	return err
}

// Update is simply a shorthand for Registry().UpdateMetric
func (w *PCPWriter) Update(m Metric, val interface{}) error { return w.Registry().UpdateMetric(m, val) }
<commit_msg>writer: make compatible with metrics and registry changes<commit_after>package speed

import (
	"errors"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/performancecopilot/speed/bytebuffer"
)

// byte lengths of different components in an mmv file
const (
	HeaderLength         = 40
	TocLength            = 16
	MetricLength         = 104
	ValueLength          = 32
	InstanceLength       = 80
	InstanceDomainLength = 32
	StringBlockLength    = 256
)

// MaxMetricNameLength is the maximum length for a metric name
const MaxMetricNameLength = 63

// MaxDataValueSize is the maximum byte length for a stored metric value, unless it is a string
const MaxDataValueSize = 16

// EraseFileOnStop if set to true, will also delete the memory mapped file
var EraseFileOnStop = false

// Writer defines the interface of a MMV file writer's properties
type Writer interface {
	// a writer must contain a registry of metrics and instance domains
	Registry() Registry

	// writes an mmv file
	Start() error

	// stops writing and cleans up
	Stop() error

	// returns the number of bytes that will be written by the current writer
	Length() int

	// adds a metric to be written
	Register(Metric) error

	// adds an instance domain to be written
	RegisterIndom(InstanceDomain) error

	// adds metric and instance domain from a string
	RegisterString(string, interface{}, MetricSemantics, MetricType, MetricUnit) error

	// update a metric
	Update(Metric, interface{}) error
}

func mmvFileLocation(name string) (string, error) {
	if strings.ContainsRune(name, os.PathSeparator) {
		return "", errors.New("name cannot have path separator")
	}

	tdir, present := config["PCP_TMP_DIR"]
	var loc string
	if present {
		loc = path.Join(rootPath, tdir)
	} else {
		loc = os.TempDir()
	}

	return path.Join(loc, "mmv", name), nil
}

// PCPClusterIDBitLength is the bit length of the cluster id
// for a set of PCP metrics
const PCPClusterIDBitLength = 12

// MMVFlag represents an enumerated type to represent mmv flag values
type MMVFlag int

// values for MMVFlag
const (
	NoPrefixFlag MMVFlag = 1 << iota
	ProcessFlag
	SentinelFlag
)

//go:generate stringer -type=MMVFlag

// PCPWriter implements a writer that can write PCP compatible MMV files
type PCPWriter struct {
	sync.Mutex
	loc       string            // absolute location of the mmv file
	clusterID uint32            // cluster identifier for the writer
	flag      MMVFlag           // write flag
	r         *PCPRegistry      // current registry
	buffer    bytebuffer.Buffer // current Buffer
}

// NewPCPWriter initializes a new PCPWriter object
func NewPCPWriter(name string, flag MMVFlag) (*PCPWriter, error) {
	fileLocation, err := mmvFileLocation(name)
	if err != nil {
		return nil, err
	}

	return &PCPWriter{
		loc:       fileLocation,
		r:         NewPCPRegistry(),
		clusterID: getHash(name, PCPClusterIDBitLength),
		flag:      flag,
		buffer:    nil,
	}, nil
}

// Registry returns a writer's registry
func (w *PCPWriter) Registry() Registry {
	return w.r
}

func (w *PCPWriter) tocCount() int {
	ans := 2

	if w.Registry().InstanceCount() > 0 {
		ans += 2
	}

	if w.Registry().StringCount() > 0 {
		ans++
	}

	return ans
}

// Length returns the byte length of data in the mmv file written by the current writer
func (w *PCPWriter) Length() int {
	return HeaderLength +
		(w.tocCount() * TocLength) +
		(w.Registry().InstanceCount() * InstanceLength) +
		(w.Registry().InstanceDomainCount() * InstanceDomainLength) +
		(w.Registry().MetricCount() * (MetricLength + ValueLength)) +
		(w.Registry().StringCount() * StringBlockLength)
}

func (w *PCPWriter) initializeOffsets() {
	indomoffset := HeaderLength + TocLength*w.tocCount()
	instanceoffset := indomoffset + InstanceDomainLength*w.r.InstanceDomainCount()
	metricsoffset := instanceoffset + InstanceLength*w.r.InstanceCount()
	valuesoffset := metricsoffset + MetricLength*w.r.MetricCount()
	stringsoffset := valuesoffset + ValueLength*w.r.MetricCount()

	w.r.indomoffset = indomoffset
	w.r.instanceoffset = instanceoffset
	w.r.metricsoffset = metricsoffset
	w.r.valuesoffset = valuesoffset
	w.r.stringsoffset = stringsoffset

	initializeSingletonMetricOffsets := func(metric *PCPSingletonMetric) {
		metric.descoffset = metricsoffset
		metricsoffset += MetricLength
		metric.valueoffset = valuesoffset
		valuesoffset += ValueLength

		if metric.shortDescription.val != "" {
			metric.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if metric.longDescription.val != "" {
			metric.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}

	for _, indom := range w.r.instanceDomains {
		indom.offset = indomoffset
		indom.instanceOffset = instanceoffset
		indomoffset += InstanceDomainLength

		for _, i := range indom.instances {
			i.descoffset = instanceoffset
			instanceoffset += InstanceLength
		}

		if indom.shortDescription.val != "" {
			indom.shortDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}

		if indom.longDescription.val != "" {
			indom.longDescription.offset = stringsoffset
			stringsoffset += StringBlockLength
		}
	}

	for _, metric := range w.r.metrics {
		switch metric.(type) {
		case *PCPSingletonMetric:
			initializeSingletonMetricOffsets(metric.(*PCPSingletonMetric))
		}
	}
}

func (w *PCPWriter) writeHeaderBlock() (generation2offset int, generation int64) {
	// tag
	w.buffer.WriteString("MMV")
	w.buffer.SetPos(w.buffer.Pos() + 1) // extra null byte is needed and \0 isn't a valid escape character in go

	// version
	w.buffer.WriteUint32(1)

	// generation
	generation = time.Now().Unix()
	w.buffer.WriteInt64(generation)

	generation2offset = w.buffer.Pos()

	w.buffer.WriteInt64(0)

	// tocCount
	w.buffer.WriteInt32(int32(w.tocCount()))

	// flag mask
	w.buffer.WriteInt32(int32(w.flag))

	// process identifier
	w.buffer.WriteInt32(int32(os.Getpid()))

	// cluster identifier
	w.buffer.WriteUint32(w.clusterID)

	return
}

func (w *PCPWriter) writeSingleToc(pos, identifier, count, offset int) {
	w.buffer.SetPos(pos)
	w.buffer.WriteInt32(int32(identifier))
	w.buffer.WriteInt32(int32(count))
	w.buffer.WriteUint64(uint64(offset))
}

func (w *PCPWriter) writeTocBlock() {
	tocpos := HeaderLength

	// instance domains toc
	if w.Registry().InstanceDomainCount() > 0 {
		// 1 is the identifier for instance domains
		w.writeSingleToc(tocpos, 1, w.r.InstanceDomainCount(), w.r.indomoffset)
		tocpos += TocLength
	}

	// instances toc
	if w.Registry().InstanceCount() > 0 {
		// 2 is the identifier for instances
		w.writeSingleToc(tocpos, 2, w.r.InstanceCount(), w.r.instanceoffset)
		tocpos += TocLength
	}

	// metrics and values toc
	metricsoffset, valuesoffset := w.r.metricsoffset, w.r.valuesoffset
	if w.Registry().MetricCount() == 0 {
		metricsoffset, valuesoffset = 0, 0
	}

	// 3 is the identifier for metrics
	w.writeSingleToc(tocpos, 3, w.r.MetricCount(), metricsoffset)
	tocpos += TocLength

	// 4 is the identifier for values
	w.writeSingleToc(tocpos, 4, w.r.MetricCount(), valuesoffset)
	tocpos += TocLength

	// strings toc
	if w.Registry().StringCount() > 0 {
		// 5 is the identifier for strings
		w.writeSingleToc(tocpos, 5, w.r.StringCount(), w.r.stringsoffset)
	}
}

func (w *PCPWriter) writeInstanceAndInstanceDomainBlock() {
	for _, indom := range w.r.instanceDomains {
		w.buffer.SetPos(indom.offset)
		w.buffer.WriteUint32(indom.ID())
		w.buffer.WriteInt32(int32(indom.InstanceCount()))
		w.buffer.WriteInt64(int64(indom.instanceOffset))

		so, lo := indom.shortDescription.offset, indom.longDescription.offset
		w.buffer.WriteInt64(int64(so))
		w.buffer.WriteInt64(int64(lo))

		if so != 0 {
			w.buffer.SetPos(so)
			w.buffer.WriteString(indom.shortDescription.val)
		}

		if lo != 0 {
			w.buffer.SetPos(lo)
			w.buffer.WriteString(indom.longDescription.val)
		}

		for _, i := range indom.instances {
			w.buffer.SetPos(i.descoffset)
			w.buffer.WriteInt64(int64(indom.offset))
			w.buffer.WriteInt32(0)
			w.buffer.WriteUint32(i.id)
			w.buffer.WriteString(i.name)
		}
	}
}

func (w *PCPWriter) writeMetricDesc(m PCPMetric, pos int) {
	w.buffer.SetPos(pos)

	w.buffer.WriteString(m.Name())
	w.buffer.SetPos(pos + MaxMetricNameLength + 1)
	w.buffer.WriteUint32(m.ID())
	w.buffer.WriteInt32(int32(m.Type()))
	w.buffer.WriteInt32(int32(m.Semantics()))
	w.buffer.WriteUint32(m.Unit().PMAPI())
	if m.Indom() != nil {
		w.buffer.WriteUint32(m.Indom().ID())
	} else {
		w.buffer.WriteInt32(-1)
	}
	w.buffer.WriteInt32(0)

	so, lo := m.ShortDescription().offset, m.LongDescription().offset
	w.buffer.WriteInt64(int64(so))
	w.buffer.WriteInt64(int64(lo))

	if so != 0 {
		w.buffer.SetPos(so)
		w.buffer.WriteString(m.ShortDescription().val)
	}

	if lo != 0 {
		w.buffer.SetPos(lo)
		w.buffer.WriteString(m.LongDescription().val)
	}
}

func (w *PCPWriter) writeSingletonMetric(m *PCPSingletonMetric) {
	w.writeMetricDesc(m, m.descoffset)

	pos := m.valueoffset
	w.buffer.SetPos(pos)

	m.t.WriteVal(m.val, w.buffer)

	w.buffer.SetPos(pos + MaxDataValueSize)
	w.buffer.WriteInt64(int64(m.descoffset))
	w.buffer.WriteInt64(0)
}

func (w *PCPWriter) writeMetricsAndValuesBlock() {
	for _, metric := range w.r.metrics {
		switch metric.(type) {
		case *PCPSingletonMetric:
			w.writeSingletonMetric(metric.(*PCPSingletonMetric))
		}
	}
}

// fillData will fill the Buffer with the mmv file
// data as long as something doesn't go wrong
func (w *PCPWriter) fillData() error {
	generation2offset, generation := w.writeHeaderBlock()
	w.writeTocBlock()
	w.writeInstanceAndInstanceDomainBlock()
	w.writeMetricsAndValuesBlock()

	w.buffer.SetPos(generation2offset)
	w.buffer.WriteUint64(uint64(generation))

	return nil
}

// Start dumps existing registry data
func (w *PCPWriter) Start() error {
	w.Lock()
	defer w.Unlock()

	l := w.Length()

	w.initializeOffsets()

	buffer, err := bytebuffer.NewMemoryMappedBuffer(w.loc, l)
	if err != nil {
		return err
	}
	w.buffer = buffer

	w.fillData()

	w.r.mapped = true

	return nil
}

// Stop removes existing mapping and cleans up
func (w *PCPWriter) Stop() error {
	w.Lock()
	defer w.Unlock()

	w.r.mapped = false

	err := w.buffer.(*bytebuffer.MemoryMappedBuffer).Unmap(EraseFileOnStop)
	w.buffer = nil
	if err != nil {
		return err
	}

	return nil
}

// Register is simply a shorthand for Registry().AddMetric
func (w *PCPWriter) Register(m Metric) error { return w.Registry().AddMetric(m) }

// RegisterIndom is simply a shorthand for Registry().AddInstanceDomain
func (w *PCPWriter) RegisterIndom(indom InstanceDomain) error {
	return w.Registry().AddInstanceDomain(indom)
}

// RegisterString is simply a shorthand for Registry().AddMetricByString
func (w *PCPWriter) RegisterString(str string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) error {
	_, err := w.Registry().AddSingletonMetricByString(str, val, s, t, u)
	return err
}
<|endoftext|>


<commit_before>package lookup

import (
	"fmt"
	"koding/kites/kloud/multiec2"
	"sync"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/ec2"
)

type Lookup struct {
	// values contains a list of instance tags that are identified as test
	// instances. By default all instances are fetched.
	values []string

	clients *multiec2.Clients
}

func NewAWS(auth aws.Auth) *Lookup {
	return &Lookup{
		clients: multiec2.New(auth, []string{
			"us-east-1",
			"ap-southeast-1",
			"us-west-2",
			"eu-west-1",
		}),
	}
}

// Instances returns all instances that belongs to the given client/region if
func (l *Lookup) Instances(client *ec2.EC2) (Instances, error) {
	instances := make([]ec2.Instance, 0)

	opts := &ec2.InstancesOpts{
		MaxResults: 900,
	}

	resp, err := client.InstancesWithOpts([]string{}, opts)
	if err != nil {
		return nil, err
	}

	for _, res := range resp.Reservations {
		instances = append(instances, res.Instances...)
	}

	nextToken := resp.NextToken

	// get all results until nextToken is empty
	for nextToken != "" {
		opts := &ec2.InstancesOpts{
			NextToken: nextToken,
		}

		resp, err := client.InstancesWithOpts([]string{}, opts)
		if err != nil {
			return nil, err
		}

		for _, res := range resp.Reservations {
			instances = append(instances, res.Instances...)
		}

		nextToken = resp.NextToken
	}

	m := make(Instances, len(instances))

	for _, instance := range instances {
		m[instance.InstanceId] = instance
	}

	return m, nil
}

// FetchInstances fetches all instances from all regions
func (l *Lookup) FetchInstances() *MultiInstances {
	var wg sync.WaitGroup

	allInstances := NewMultiInstanes(l.clients)

	for region, client := range l.clients.Regions() {
		wg.Add(1)
		go func(region string, client *ec2.EC2) {
			defer wg.Done()

			instances, err := l.Instances(client)
			if err != nil {
				fmt.Printf("[%s] fetching error: %s\n", region, err)
				return
			}

			allInstances.Add(client, instances)
		}(region, client)
	}

	wg.Wait()

	return allInstances
}
<commit_msg>cleaner: lower to avoid timeouts<commit_after>package lookup

import (
	"fmt"
	"koding/kites/kloud/multiec2"
	"sync"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/ec2"
)

type Lookup struct {
	// values contains a list of instance tags that are identified as test
	// instances. By default all instances are fetched.
	values []string

	clients *multiec2.Clients
}

func NewAWS(auth aws.Auth) *Lookup {
	return &Lookup{
		clients: multiec2.New(auth, []string{
			"us-east-1",
			"ap-southeast-1",
			"us-west-2",
			"eu-west-1",
		}),
	}
}

// Instances returns all instances that belongs to the given client/region if
func (l *Lookup) Instances(client *ec2.EC2) (Instances, error) {
	instances := make([]ec2.Instance, 0)

	opts := &ec2.InstancesOpts{
		MaxResults: 500,
	}

	resp, err := client.InstancesWithOpts([]string{}, opts)
	if err != nil {
		return nil, err
	}

	for _, res := range resp.Reservations {
		instances = append(instances, res.Instances...)
	}

	nextToken := resp.NextToken

	// get all results until nextToken is empty
	for nextToken != "" {
		opts := &ec2.InstancesOpts{
			NextToken: nextToken,
		}

		resp, err := client.InstancesWithOpts([]string{}, opts)
		if err != nil {
			return nil, err
		}

		for _, res := range resp.Reservations {
			instances = append(instances, res.Instances...)
		}

		nextToken = resp.NextToken
	}

	m := make(Instances, len(instances))

	for _, instance := range instances {
		m[instance.InstanceId] = instance
	}

	return m, nil
}

// FetchInstances fetches all instances from all regions
func (l *Lookup) FetchInstances() *MultiInstances {
	var wg sync.WaitGroup

	allInstances := NewMultiInstanes(l.clients)

	for region, client := range l.clients.Regions() {
		wg.Add(1)
		go func(region string, client *ec2.EC2) {
			defer wg.Done()

			instances, err := l.Instances(client)
			if err != nil {
				fmt.Printf("[%s] fetching error: %s\n", region, err)
				return
			}

			allInstances.Add(client, instances)
		}(region, client)
	}

	wg.Wait()

	return allInstances
}
<|endoftext|>


<commit_before>package db

import (
	"github.com/jinzhu/gorm"

	"github.com/satori/go.uuid"
)

func (u UserModel) UpdateRegistrationID(database *gorm.DB, newRegistrationID *uuid.UUID) error {
	return database.Model(&u).Update("registration_id", newRegistrationID).Error
}

func SelectUserByLogin(database *gorm.DB, login string) (UserModel, error) {

	var user = UserModel{}
	err := database.Where("login = $1", login).First(&user).Error

	return user, err
}

func SelectUserByRegisterID(database *gorm.DB, registerID uuid.UUID) (UserModel, error) {
	var user = UserModel{}
	err := database.Where("registration_id = $1", registerID).First(&user).Error

	return user, err
}

func GetUserPasswordHash(database *gorm.DB, login string) ([]byte, error) {
	var model, err = SelectUserByLogin(database, login)

	return model.PasswordHash, err
}

func IsLoginExists(database *gorm.DB, login string) (bool, error) {
	var exists bool
	var err = database.Raw(existsLoginValidationSQL, login).Scan(&exists).Error

	return exists, err
}
<commit_msg>fix raw<commit_after>package db

import (
	"github.com/jinzhu/gorm"

	"github.com/satori/go.uuid"
)

func (u UserModel) UpdateRegistrationID(database *gorm.DB, newRegistrationID *uuid.UUID) error {
	return database.Model(&u).Update("registration_id", newRegistrationID).Error
}

func SelectUserByLogin(database *gorm.DB, login string) (UserModel, error) {

	var user = UserModel{}
	err := database.Where("login = $1", login).First(&user).Error

	return user, err
}

func SelectUserByRegisterID(database *gorm.DB, registerID uuid.UUID) (UserModel, error) {
	var user = UserModel{}
	err := database.Where("registration_id = $1", registerID).First(&user).Error

	return user, err
}

func GetUserPasswordHash(database *gorm.DB, login string) ([]byte, error) {
	var model, err = SelectUserByLogin(database, login)

	return model.PasswordHash, err
}

func IsLoginExists(database *gorm.DB, login string) (bool, error) {
	var exists bool
	var err = database.DB().QueryRow(existsLoginValidationSQL, login).Scan(&exists)

	return exists, err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/lib/pq"
)

func sqlInterface(
	rootAccount *Account,
	transactions []*Transaction,
) {

	dbDir := filepath.Join(os.TempDir(), fmt.Sprintf("keep-%d", rand.Int63()))
	if out, err := exec.Command("initdb", "-D", dbDir).CombinedOutput(); err != nil {
		fail("%v: %s", err, out)
	}
	pt("db dir: %s\n", dbDir)
	defer exec.Command("rm", "-rf", dbDir).Run()

	port := 10000 + rand.Intn(50000)
	conf := `
		port = ` + fmt.Sprintf("%d", port) + `
	`
	if err := ioutil.WriteFile(filepath.Join(dbDir, "postgresql.conf"), []byte(conf), 0644); err != nil {
		fail("%v", err)
	}
	c := exec.Command("postgres", "-D", dbDir)
	c.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}
	if err := c.Start(); err != nil {
		fail("%v", err)
	}
	defer syscall.Kill(-c.Process.Pid, syscall.SIGKILL)
	time.Sleep(time.Second)
	pt("db started\n")

	db, err := sqlx.Open("postgres", fmt.Sprintf("postgres://localhost:%d/postgres?sslmode=disable", port))
	if err != nil {
		fail("%v", err)
	}
	defer db.Close()
	if _, err := db.Exec(`
		CREATE TABLE entries (
			id bigserial primary key,
			transaction bigint,
			transaction_description text,
			time timestamp with time zone,
			account text[],
			currency text,
			amount numeric,
			description text
		)
		`,
	); err != nil {
		fail("%v", err)
	}
	tx := db.MustBegin()
	for tid, transaction := range transactions {
		for _, entry := range transaction.Entries {
			if _, err := tx.Exec(`
				INSERT INTO entries
				(
					transaction, transaction_description,
					time, account, currency, amount, description
				)
				VALUES (
					$1, $2,
					$3, $4, $5, $6::numeric, $7
				)
				`,
				tid,
				transaction.Description,
				entry.Time,
				func() (ret pq.StringArray) {
					acc := entry.Account
					for acc != rootAccount {
						ret = append(ret, acc.Name)
						acc = acc.Parent
					}
					for i := len(ret)/2 - 1; i >= 0; i-- {
						j := len(ret) - 1 - i
						ret[i], ret[j] = ret[j], ret[i]
					}
					return
				}(),
				entry.Currency,
				entry.Amount.FloatString(3),
				entry.Description,
			); err != nil {
				fail("%v", err)
			}
		}
	}
	if err := tx.Commit(); err != nil {
		fail("%v", err)
	}
	pt("data loaded\n")

	psql := exec.Command("psql", fmt.Sprintf("postgres://localhost:%d/postgres", port))
	psql.Stdout = os.Stdout
	psql.Stdin = os.Stdin
	psql.Stderr = os.Stderr
	if err := psql.Run(); err != nil {
		fail("%v", err)
	}
}
<commit_msg>add SQL views<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/lib/pq"
)

func sqlInterface(
	rootAccount *Account,
	transactions []*Transaction,
) {

	dbDir := filepath.Join(os.TempDir(), fmt.Sprintf("keep-%d", rand.Int63()))
	if out, err := exec.Command("initdb", "-D", dbDir).CombinedOutput(); err != nil {
		fail("%v: %s", err, out)
	}
	pt("db dir: %s\n", dbDir)
	defer exec.Command("rm", "-rf", dbDir).Run()

	port := 10000 + rand.Intn(50000)
	conf := `
		port = ` + fmt.Sprintf("%d", port) + `
	`
	if err := ioutil.WriteFile(filepath.Join(dbDir, "postgresql.conf"), []byte(conf), 0644); err != nil {
		fail("%v", err)
	}
	c := exec.Command("postgres", "-D", dbDir)
	c.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}
	if err := c.Start(); err != nil {
		fail("%v", err)
	}
	defer syscall.Kill(-c.Process.Pid, syscall.SIGKILL)
	time.Sleep(time.Second)
	pt("db started\n")

	db, err := sqlx.Open("postgres", fmt.Sprintf("postgres://localhost:%d/postgres?sslmode=disable", port))
	if err != nil {
		fail("%v", err)
	}
	defer db.Close()
	tx := db.MustBegin()
	if _, err := tx.Exec(`
		CREATE TABLE entries (
			id bigserial primary key,
			transaction bigint,
			transaction_description text,
			time timestamp with time zone,
			account text[],
			currency text,
			amount numeric,
			description text
		)
		`,
	); err != nil {
		fail("%v", err)
	}
	for _, view := range views {
		if _, err := tx.Exec(view); err != nil {
			fail("%v", err)
		}
	}
	for tid, transaction := range transactions {
		for _, entry := range transaction.Entries {
			if _, err := tx.Exec(`
				INSERT INTO entries
				(
					transaction, transaction_description,
					time, account, currency, amount, description
				)
				VALUES (
					$1, $2,
					$3, $4, $5, $6::numeric, $7
				)
				`,
				tid,
				transaction.Description,
				entry.Time,
				func() (ret pq.StringArray) {
					acc := entry.Account
					for acc != rootAccount {
						ret = append(ret, acc.Name)
						acc = acc.Parent
					}
					for i := len(ret)/2 - 1; i >= 0; i-- {
						j := len(ret) - 1 - i
						ret[i], ret[j] = ret[j], ret[i]
					}
					return
				}(),
				entry.Currency,
				entry.Amount.FloatString(3),
				entry.Description,
			); err != nil {
				fail("%v", err)
			}
		}
	}
	if err := tx.Commit(); err != nil {
		fail("%v", err)
	}
	pt("data loaded\n")

	sigs := make(chan os.Signal)
	go func() {
		for {
			<-sigs
		}
	}()
	signal.Notify(sigs, os.Interrupt)

	psql := exec.Command("psql", fmt.Sprintf("postgres://localhost:%d/postgres", port))
	psql.Stdout = os.Stdout
	psql.Stdin = os.Stdin
	psql.Stderr = os.Stderr
	if err := psql.Run(); err != nil {
		fail("%v", err)
	}
}

var views = []string{
	// props
	`
	create view props as
	select distinct on (time, transaction)
	time, transaction_description
	from entries
	where account[1] = '支出'
	and account[2] in (
		'数码',
		'物品',
		'衣物服饰',
		'消耗品',
		'保健品',
		'书籍',
		'药物',
		'性用品'
	)
	order by time desc, transaction
	`,

	// monthly_expense
	`
	create view monthly_expense as
	select
	to_char(date_trunc('month', time), 'YYYY-MM') as month,
	currency,
	sum(amount),
	(
		select string_agg(
			account || currency || amount::text,
			E'\n'
		) from (
			select account[2] as account, currency, sum(amount) as amount
			from unnest(array_agg(id)) as id
			join entries e2 using (id)
			group by account[2], currency
			order by amount desc
		) t0
	)
	from entries
	where account[1] = '支出'
	group by date_trunc('month', time), currency
	order by month desc, currency asc
	`,

	// monthly_income
	`
	create view monthly_income as
	select
	to_char(date_trunc('month', time), 'YYYY-MM') as month,
	currency,
	sum(amount),
	(
		select string_agg(
			account || currency || amount::text,
			E'\n'
		) from (
			select account[2] as account, currency, sum(amount) as amount
			from unnest(array_agg(id)) as id
			join entries e2 using (id)
			group by account[2], currency
			order by amount asc
		) t0
	)
	from entries
	where account[1] = '收入'
	group by date_trunc('month', time), currency
	order by month desc, currency asc
	`,

	//
}
<|endoftext|>


<commit_before>// Copyright 2015 - António Meireles  <antonio.meireles@reformi.st>
//
// 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 main

import (
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"github.com/rakyll/pb"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

var (
	sshCmd = &cobra.Command{
		Use:     "ssh",
		Aliases: []string{"attach"},
		Short:   "Attach to or run commands inside a running CoreOS instance",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			vipre.BindPFlags(cmd.Flags())
			if len(args) < 1 {
				return fmt.Errorf("This command requires at least " +
					"one argument to work ")
			}
			return
		},
		RunE: sshCommand,
		Example: `  corectl ssh VMid                 // logins into VMid
  corectl ssh VMid "some commands" // runs 'some commands' inside VMid and exits`,
	}
	scpCmd = &cobra.Command{
		Use:     "put",
		Aliases: []string{"copy", "cp", "scp"},
		Short:   "copy file to inside VM",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			vipre.BindPFlags(cmd.Flags())
			if len(args) < 2 {
				return fmt.Errorf("This command requires at least " +
					"two argument to work ")
			}
			return
		},
		RunE: scpCommand,
		Example: `  // copies 'filePath' into '/destinationPath' inside VMid
  corectl put filePath VMid:/destinationPath`,
	}
)

func sshCommand(cmd *cobra.Command, args []string) (err error) {
	var sshSession = &sshClient{}
	vm := VMInfo{}

	if vm, err = vmInfo(args[0]); err != nil {
		return
	}

	if sshSession, err = vm.startSSHsession(); err != nil {
		return
	}
	defer sshSession.close()

	if len(args) > 1 {
		return sshSession.executeRemoteCommand(strings.Join(args[1:], " "))
	}
	return sshSession.remoteShell()
}

type sshClient struct {
	session                   *ssh.Session
	conn                      *ssh.Client
	oldState                  *terminal.State
	termWidth, termHeight, fd int
}

func (c *sshClient) close() {
	c.conn.Close()
	c.session.Close()
	terminal.Restore(c.fd, c.oldState)
}

func (vm VMInfo) startSSHsession() (c *sshClient, err error) {
	var secret ssh.Signer
	c = &sshClient{}

	if secret, err = ssh.ParsePrivateKey(
		[]byte(vm.InternalSSHprivKey)); err != nil {
		return
	}

	config := &ssh.ClientConfig{
		User: "core", Auth: []ssh.AuthMethod{
			ssh.PublicKeys(secret),
		},
	}

	//wait a bit for VM's ssh to be available...
	for {
		var e error
		if c.conn, e = ssh.Dial("tcp", vm.PublicIP+":22", config); e == nil {
			break
		}
		time.Sleep(100 * time.Millisecond)
		select {
		case <-time.After(time.Second * 5):
			return c, fmt.Errorf("%s unreachable", vm.PublicIP+":22")
		}
	}

	if c.session, err = c.conn.NewSession(); err != nil {
		return c, fmt.Errorf("unable to create session: %s", err)
	}

	c.fd = int(os.Stdin.Fd())
	if c.oldState, err = terminal.MakeRaw(c.fd); err != nil {
		return
	}

	c.session.Stdout, c.session.Stderr, c.session.Stdin =
		os.Stdout, os.Stderr, os.Stdin

	if c.termWidth, c.termHeight, err = terminal.GetSize(c.fd); err != nil {
		return
	}

	modes := ssh.TerminalModes{
		ssh.ECHO: 1, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400,
	}

	if err = c.session.RequestPty("xterm-256color",
		c.termHeight, c.termWidth, modes); err != nil {
		return c, fmt.Errorf("request for pseudo terminal failed: %s", err)
	}
	return
}

func (c *sshClient) executeRemoteCommand(run string) (err error) {
	if err = c.session.Run(run); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return
	}
	return nil
}

func (c *sshClient) remoteShell() (err error) {
	if err = c.session.Shell(); err != nil {
		return
	}

	if err = c.session.Wait(); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return err
	}
	return nil
}

func vmInfo(id string) (vm VMInfo, err error) {
	var up []VMInfo
	if up, err = allRunningInstances(); err != nil {
		return
	}
	for _, v := range up {
		if v.Name == id || v.UUID == id {
			return v, err
		}
	}
	return vm, fmt.Errorf("'%s' not found, or dead", id)
}

func (c *sshClient) sCopy(source, destination, target string) (err error) {
	var (
		ftp         *sftp.Client
		src         *os.File
		srcS, destS os.FileInfo
		dest        *sftp.File
		bar         *pb.ProgressBar
	)

	if ftp, err = sftp.NewClient(c.conn); err != nil {
		return
	}
	defer ftp.Close()
	log.Println("uploading '" + source + "' to '" +
		target + ":" + destination + "'")

	if src, err = os.Open(source); err != nil {
		return
	}
	defer src.Close()
	if srcS, err = os.Stat(source); err != nil {
		return
	}

	if dest, err = ftp.Create(destination); err != nil {
        // XXX
		return
	}
	defer dest.Close()
	bar = pb.New(int(srcS.Size())).SetUnits(pb.U_BYTES)
	bar.Start()
	defer bar.Finish()
	if _, err = io.Copy(dest, src); err != nil {
		return
	}

	if destS, err = ftp.Stat(destination); err != nil {
		return
	}
	if srcS.Size() != destS.Size() {
		err = fmt.Errorf("something went wrong. " +
			"destination file size != from sources'")
	}
	return
}

func scpCommand(cmd *cobra.Command, args []string) (err error) {
	var (
		session, vm                 = &sshClient{}, VMInfo{}
		split                       = strings.Split(args[1], ":")
		source, destination, target = args[0], split[1], split[0]
	)
	if vm, err = vmInfo(target); err != nil {
		return
	}
	if session, err = vm.startSSHsession(); err != nil {
		return
	}
	defer session.close()
	return session.sCopy(source, destination, target)
}

func init() {
	RootCmd.AddCommand(sshCmd)
	RootCmd.AddCommand(scpCmd)
}
<commit_msg>checks that destination dir exists on target before attempting to upload.<commit_after>// Copyright 2015 - António Meireles  <antonio.meireles@reformi.st>
//
// 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 main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"github.com/rakyll/pb"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

var (
	sshCmd = &cobra.Command{
		Use:     "ssh",
		Aliases: []string{"attach"},
		Short:   "Attach to or run commands inside a running CoreOS instance",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			vipre.BindPFlags(cmd.Flags())
			if len(args) < 1 {
				return fmt.Errorf("This command requires at least " +
					"one argument to work ")
			}
			return
		},
		RunE: sshCommand,
		Example: `  corectl ssh VMid                 // logins into VMid
  corectl ssh VMid "some commands" // runs 'some commands' inside VMid and exits`,
	}
	scpCmd = &cobra.Command{
		Use:     "put",
		Aliases: []string{"copy", "cp", "scp"},
		Short:   "copy file to inside VM",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			vipre.BindPFlags(cmd.Flags())
			if len(args) < 2 {
				return fmt.Errorf("This command requires at least " +
					"two argument to work ")
			}
			return
		},
		RunE: scpCommand,
		Example: `  // copies 'filePath' into '/destinationPath' inside VMid
  corectl put filePath VMid:/destinationPath`,
	}
)

func sshCommand(cmd *cobra.Command, args []string) (err error) {
	var sshSession = &sshClient{}
	vm := VMInfo{}

	if vm, err = vmInfo(args[0]); err != nil {
		return
	}

	if sshSession, err = vm.startSSHsession(); err != nil {
		return
	}
	defer sshSession.close()

	if len(args) > 1 {
		return sshSession.executeRemoteCommand(strings.Join(args[1:], " "))
	}
	return sshSession.remoteShell()
}

type sshClient struct {
	session                   *ssh.Session
	conn                      *ssh.Client
	oldState                  *terminal.State
	termWidth, termHeight, fd int
}

func (c *sshClient) close() {
	c.conn.Close()
	c.session.Close()
	terminal.Restore(c.fd, c.oldState)
}

func (vm VMInfo) startSSHsession() (c *sshClient, err error) {
	var secret ssh.Signer
	c = &sshClient{}

	if secret, err = ssh.ParsePrivateKey(
		[]byte(vm.InternalSSHprivKey)); err != nil {
		return
	}

	config := &ssh.ClientConfig{
		User: "core", Auth: []ssh.AuthMethod{
			ssh.PublicKeys(secret),
		},
	}

	//wait a bit for VM's ssh to be available...
	for {
		var e error
		if c.conn, e = ssh.Dial("tcp", vm.PublicIP+":22", config); e == nil {
			break
		}
		time.Sleep(100 * time.Millisecond)
		select {
		case <-time.After(time.Second * 5):
			return c, fmt.Errorf("%s unreachable", vm.PublicIP+":22")
		}
	}

	if c.session, err = c.conn.NewSession(); err != nil {
		return c, fmt.Errorf("unable to create session: %s", err)
	}

	c.fd = int(os.Stdin.Fd())
	if c.oldState, err = terminal.MakeRaw(c.fd); err != nil {
		return
	}

	c.session.Stdout, c.session.Stderr, c.session.Stdin =
		os.Stdout, os.Stderr, os.Stdin

	if c.termWidth, c.termHeight, err = terminal.GetSize(c.fd); err != nil {
		return
	}

	modes := ssh.TerminalModes{
		ssh.ECHO: 1, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400,
	}

	if err = c.session.RequestPty("xterm-256color",
		c.termHeight, c.termWidth, modes); err != nil {
		return c, fmt.Errorf("request for pseudo terminal failed: %s", err)
	}
	return
}

func (c *sshClient) executeRemoteCommand(run string) (err error) {
	if err = c.session.Run(run); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return
	}
	return nil
}

func (c *sshClient) remoteShell() (err error) {
	if err = c.session.Shell(); err != nil {
		return
	}

	if err = c.session.Wait(); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return err
	}
	return nil
}

func vmInfo(id string) (vm VMInfo, err error) {
	var up []VMInfo
	if up, err = allRunningInstances(); err != nil {
		return
	}
	for _, v := range up {
		if v.Name == id || v.UUID == id {
			return v, err
		}
	}
	return vm, fmt.Errorf("'%s' not found, or dead", id)
}

func (c *sshClient) sCopy(source, destination, target string) (err error) {
	var (
		ftp         *sftp.Client
		src         *os.File
		srcS, destS os.FileInfo
		dest        *sftp.File
		bar         *pb.ProgressBar
	)

	if ftp, err = sftp.NewClient(c.conn); err != nil {
		return
	}
	defer ftp.Close()
	log.Println("uploading '" + source + "' to '" +
		target + ":" + destination + "'")

	if src, err = os.Open(source); err != nil {
		return
	}
	defer src.Close()
	if srcS, err = os.Stat(source); err != nil {
		return
	}
	if _, err = ftp.ReadDir(filepath.Dir(destination)); err != nil {
		err = fmt.Errorf("unable to upload %v as parent %v "+
			"not in target", source, filepath.Dir(destination))
		return
	}
	if dest, err = ftp.Create(destination); err != nil {
		return
	}
	defer dest.Close()
	bar = pb.New(int(srcS.Size())).SetUnits(pb.U_BYTES)
	bar.Start()
	defer bar.Finish()
	if _, err = io.Copy(dest, src); err != nil {
		return
	}

	if destS, err = ftp.Stat(destination); err != nil {
		return
	}
	if srcS.Size() != destS.Size() {
		err = fmt.Errorf("something went wrong. " +
			"destination file size != from sources'")
	}
	return
}

func scpCommand(cmd *cobra.Command, args []string) (err error) {
	var (
		session, vm                 = &sshClient{}, VMInfo{}
		split                       = strings.Split(args[1], ":")
		source, destination, target = args[0], split[1], split[0]
	)
	if vm, err = vmInfo(target); err != nil {
		return
	}
	if session, err = vm.startSSHsession(); err != nil {
		return
	}
	defer session.close()
	return session.sCopy(source, destination, target)
}

func init() {
	RootCmd.AddCommand(sshCmd)
	RootCmd.AddCommand(scpCmd)
}
<|endoftext|>


<commit_before>// Copyright 2015 - António Meireles  <antonio.meireles@reformi.st>
//
// 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 main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"github.com/rakyll/pb"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

var (
	sshCmd = &cobra.Command{
		Use:     "ssh",
		Aliases: []string{"attach"},
		Short:   "Attach to or run commands inside a running CoreOS instance",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			engine.rawArgs.BindPFlags(cmd.Flags())
			if len(args) < 1 {
				return fmt.Errorf("This command requires at least " +
					"one argument to work ")
			}
			return
		},
		RunE: sshCommand,
		Example: `  corectl ssh VMid                 // logins into VMid
  corectl ssh VMid "some commands" // runs 'some commands' inside VMid and exits`,
	}
	scpCmd = &cobra.Command{
		Use:     "put",
		Aliases: []string{"copy", "cp", "scp"},
		Short:   "copy file to inside VM",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			engine.rawArgs.BindPFlags(cmd.Flags())
			if len(args) < 2 {
				return fmt.Errorf("This command requires at least " +
					"two argument to work ")
			}
			return
		},
		RunE: scpCommand,
		Example: `  // copies 'filePath' into '/destinationPath' inside VMid
  corectl put filePath VMid:/destinationPath`,
	}
)

func sshCommand(cmd *cobra.Command, args []string) (err error) {
	var sshSession = &sshClient{}
	vm := VMInfo{}

	if vm, err = vmInfo(args[0]); err != nil {
		return
	}

	if sshSession, err = vm.startSSHsession(); err != nil {
		return
	}
	defer sshSession.close()

	if len(args) > 1 {
		return sshSession.executeRemoteCommand(strings.Join(args[1:], " "))
	}
	return sshSession.remoteShell()
}

type sshClient struct {
	session                   *ssh.Session
	conn                      *ssh.Client
	oldState                  *terminal.State
	termWidth, termHeight, fd int
}

func (c *sshClient) close() {
	c.conn.Close()
	c.session.Close()
	terminal.Restore(c.fd, c.oldState)
}

func (vm VMInfo) startSSHsession() (c *sshClient, err error) {
	var secret ssh.Signer
	c = &sshClient{}

	if secret, err = ssh.ParsePrivateKey(
		[]byte(vm.InternalSSHprivKey)); err != nil {
		return
	}

	config := &ssh.ClientConfig{
		User: "core", Auth: []ssh.AuthMethod{
			ssh.PublicKeys(secret),
		},
	}

	//wait a bit for VM's ssh to be available...
	for {
		var e error
		if c.conn, e = ssh.Dial("tcp", vm.PublicIP+":22", config); e == nil {
			break
		}
		time.Sleep(100 * time.Millisecond)
		select {
		case <-time.After(time.Second * 5):
			return c, fmt.Errorf("%s unreachable", vm.PublicIP+":22")
		}
	}

	if c.session, err = c.conn.NewSession(); err != nil {
		return c, fmt.Errorf("unable to create session: %s", err)
	}

	c.fd = int(os.Stdin.Fd())
	if c.oldState, err = terminal.MakeRaw(c.fd); err != nil {
		return
	}

	c.session.Stdout, c.session.Stderr, c.session.Stdin =
		os.Stdout, os.Stderr, os.Stdin

	if c.termWidth, c.termHeight, err = terminal.GetSize(c.fd); err != nil {
		return
	}

	modes := ssh.TerminalModes{
		ssh.ECHO: 1, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400,
	}

	if err = c.session.RequestPty("xterm-256color",
		c.termHeight, c.termWidth, modes); err != nil {
		return c, fmt.Errorf("request for pseudo terminal failed: %s", err)
	}
	return
}

func (c *sshClient) executeRemoteCommand(run string) (err error) {
	if err = c.session.Run(run); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return
	}
	return nil
}

func (c *sshClient) remoteShell() (err error) {
	if err = c.session.Shell(); err != nil {
		return
	}

	if err = c.session.Wait(); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return err
	}
	return nil
}

func vmInfo(id string) (vm VMInfo, err error) {
	var up []VMInfo
	if up, err = allRunningInstances(); err != nil {
		return
	}
	for _, v := range up {
		if v.Name == id || v.UUID == id {
			return v, err
		}
	}
	return vm, fmt.Errorf("'%s' not found, or dead", id)
}

func (c *sshClient) sCopy(source, destination, target string) (err error) {
	var (
		ftp         *sftp.Client
		src         *os.File
		srcS, destS os.FileInfo
		dest        *sftp.File
		bar         *pb.ProgressBar
	)

	if ftp, err = sftp.NewClient(c.conn); err != nil {
		return
	}
	defer ftp.Close()
	log.Println("uploading '" + source + "' to '" +
		target + ":" + destination + "'")

	if src, err = os.Open(source); err != nil {
		return
	}
	defer src.Close()
	if srcS, err = os.Stat(source); err != nil {
		return
	}
	if _, err = ftp.ReadDir(filepath.Dir(destination)); err != nil {
		err = fmt.Errorf("unable to upload %v as parent %v "+
			"not in target", source, filepath.Dir(destination))
		return
	}
	if dest, err = ftp.Create(destination); err != nil {
		return
	}
	defer dest.Close()
	bar = pb.New(int(srcS.Size())).SetUnits(pb.U_BYTES)
	bar.Start()
	defer bar.Finish()
	if _, err = io.Copy(dest, src); err != nil {
		return
	}

	if destS, err = ftp.Stat(destination); err != nil {
		return
	}
	if srcS.Size() != destS.Size() {
		err = fmt.Errorf("something went wrong. " +
			"destination file size != from sources'")
	}
	return
}

func scpCommand(cmd *cobra.Command, args []string) (err error) {
	var (
		session, vm                 = &sshClient{}, VMInfo{}
		split                       = strings.Split(args[1], ":")
		source, destination, target = args[0], split[1], split[0]
	)
	if vm, err = vmInfo(target); err != nil {
		return
	}
	if session, err = vm.startSSHsession(); err != nil {
		return
	}
	defer session.close()
	return session.sCopy(source, destination, target)
}

func init() {
	RootCmd.AddCommand(sshCmd)
	RootCmd.AddCommand(scpCmd)
}
<commit_msg>scp: handle case when destination is a directory.<commit_after>// Copyright 2015 - António Meireles  <antonio.meireles@reformi.st>
//
// 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 main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"github.com/rakyll/pb"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

var (
	sshCmd = &cobra.Command{
		Use:     "ssh",
		Aliases: []string{"attach"},
		Short:   "Attach to or run commands inside a running CoreOS instance",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			engine.rawArgs.BindPFlags(cmd.Flags())
			if len(args) < 1 {
				return fmt.Errorf("This command requires at least " +
					"one argument to work ")
			}
			return
		},
		RunE: sshCommand,
		Example: `  corectl ssh VMid                 // logins into VMid
  corectl ssh VMid "some commands" // runs 'some commands' inside VMid and exits`,
	}
	scpCmd = &cobra.Command{
		Use:     "put",
		Aliases: []string{"copy", "cp", "scp"},
		Short:   "copy file to inside VM",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			engine.rawArgs.BindPFlags(cmd.Flags())
			if len(args) < 2 {
				return fmt.Errorf("This command requires at least " +
					"two argument to work ")
			}
			return
		},
		RunE: scpCommand,
		Example: `  // copies 'filePath' into '/destinationPath' inside VMid
  corectl put filePath VMid:/destinationPath`,
	}
)

func sshCommand(cmd *cobra.Command, args []string) (err error) {
	var sshSession = &sshClient{}
	vm := VMInfo{}

	if vm, err = vmInfo(args[0]); err != nil {
		return
	}

	if sshSession, err = vm.startSSHsession(); err != nil {
		return
	}
	defer sshSession.close()

	if len(args) > 1 {
		return sshSession.executeRemoteCommand(strings.Join(args[1:], " "))
	}
	return sshSession.remoteShell()
}

type sshClient struct {
	session                   *ssh.Session
	conn                      *ssh.Client
	oldState                  *terminal.State
	termWidth, termHeight, fd int
}

func (c *sshClient) close() {
	c.conn.Close()
	c.session.Close()
	terminal.Restore(c.fd, c.oldState)
}

func (vm VMInfo) startSSHsession() (c *sshClient, err error) {
	var secret ssh.Signer
	c = &sshClient{}

	if secret, err = ssh.ParsePrivateKey(
		[]byte(vm.InternalSSHprivKey)); err != nil {
		return
	}

	config := &ssh.ClientConfig{
		User: "core", Auth: []ssh.AuthMethod{
			ssh.PublicKeys(secret),
		},
	}

	//wait a bit for VM's ssh to be available...
	for {
		var e error
		if c.conn, e = ssh.Dial("tcp", vm.PublicIP+":22", config); e == nil {
			break
		}
		time.Sleep(100 * time.Millisecond)
		select {
		case <-time.After(time.Second * 5):
			return c, fmt.Errorf("%s unreachable", vm.PublicIP+":22")
		}
	}

	if c.session, err = c.conn.NewSession(); err != nil {
		return c, fmt.Errorf("unable to create session: %s", err)
	}

	c.fd = int(os.Stdin.Fd())
	if c.oldState, err = terminal.MakeRaw(c.fd); err != nil {
		return
	}

	c.session.Stdout, c.session.Stderr, c.session.Stdin =
		os.Stdout, os.Stderr, os.Stdin

	if c.termWidth, c.termHeight, err = terminal.GetSize(c.fd); err != nil {
		return
	}

	modes := ssh.TerminalModes{
		ssh.ECHO: 1, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400,
	}

	if err = c.session.RequestPty("xterm-256color",
		c.termHeight, c.termWidth, modes); err != nil {
		return c, fmt.Errorf("request for pseudo terminal failed: %s", err)
	}
	return
}

func (c *sshClient) executeRemoteCommand(run string) (err error) {
	if err = c.session.Run(run); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return
	}
	return nil
}

func (c *sshClient) remoteShell() (err error) {
	if err = c.session.Shell(); err != nil {
		return
	}

	if err = c.session.Wait(); err != nil && !strings.HasSuffix(err.Error(),
		"exited without exit status or exit signal") {
		return err
	}
	return nil
}

func vmInfo(id string) (vm VMInfo, err error) {
	var up []VMInfo
	if up, err = allRunningInstances(); err != nil {
		return
	}
	for _, v := range up {
		if v.Name == id || v.UUID == id {
			return v, err
		}
	}
	return vm, fmt.Errorf("'%s' not found, or dead", id)
}

func (c *sshClient) sCopy(source, destination, target string) (err error) {
	var (
		ftp         *sftp.Client
		src         *os.File
		srcS, destS os.FileInfo
		dest        *sftp.File
		bar         *pb.ProgressBar
	)

	if ftp, err = sftp.NewClient(c.conn); err != nil {
		return
	}
	defer ftp.Close()

	if src, err = os.Open(source); err != nil {
		return
	}
	defer src.Close()
	if srcS, err = os.Stat(source); err != nil {
		return
	}
	if _, err = ftp.ReadDir(filepath.Dir(destination)); err != nil {
		err = fmt.Errorf("unable to upload %v as parent %v "+
			"not in target", source, filepath.Dir(destination))
		return
	}
	if _, err = ftp.ReadDir(destination); err == nil {
		destination = ftp.Join(destination, "/", filepath.Base(source))
	}
	if dest, err = ftp.Create(destination); err != nil {
		return
	}
	defer dest.Close()
	log.Println("uploading '" + source + "' to '" +
		target + ":" + destination + "'")
	bar = pb.New(int(srcS.Size())).SetUnits(pb.U_BYTES)
	bar.Start()
	defer bar.Finish()
	if _, err = io.Copy(dest, src); err != nil {
		return
	}

	if destS, err = ftp.Stat(destination); err != nil {
		return
	}
	if srcS.Size() != destS.Size() {
		err = fmt.Errorf("something went wrong. " +
			"destination file size != from sources'")
	}
	return
}

func scpCommand(cmd *cobra.Command, args []string) (err error) {
	var (
		session, vm                 = &sshClient{}, VMInfo{}
		split                       = strings.Split(args[1], ":")
		source, destination, target = args[0], split[1], split[0]
	)
	if vm, err = vmInfo(target); err != nil {
		return
	}
	if session, err = vm.startSSHsession(); err != nil {
		return
	}
	defer session.close()
	return session.sCopy(source, destination, target)
}

func init() {
	RootCmd.AddCommand(sshCmd)
	RootCmd.AddCommand(scpCmd)
}
<|endoftext|>