<commit_before>package mpredis

import (
	"flag"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/fzzy/radix/redis"
	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.redis")

// RedisPlugin mackerel plugin for Redis
type RedisPlugin struct {
	Host     string
	Port     string
	Password string
	Socket   string
	Prefix   string
	Timeout  int
	Tempfile string
}

func authenticateByPassword(c *redis.Client, password string) error {
	if r := c.Cmd("AUTH", password); r.Err != nil {
		logger.Errorf("Faild to authenticate. %s", r.Err)
		return r.Err
	}
	return nil
}

func fetchPercentageOfMemory(c *redis.Client, stat map[string]interface{}) error {
	r := c.Cmd("CONFIG", "GET", "maxmemory")
	if r.Err != nil {
		logger.Errorf("Failed to run `CONFIG GET maxmemory` command. %s", r.Err)
		return r.Err
	}

	res, err := r.Hash()
	if err != nil {
		logger.Errorf("Failed to fetch maxmemory. %s", err)
		return err
	}

	maxsize, err := strconv.ParseFloat(res["maxmemory"], 64)
	if err != nil {
		logger.Errorf("Failed to parse maxmemory. %s", err)
		return err
	}

	if maxsize == 0.0 {
		stat["percentage_of_memory"] = 0.0
	} else {
		stat["percentage_of_memory"] = 100.0 * stat["used_memory"].(float64) / maxsize
	}

	return nil
}

func fetchPercentageOfClients(c *redis.Client, stat map[string]interface{}) error {
	r := c.Cmd("CONFIG", "GET", "maxclients")
	if r.Err != nil {
		logger.Errorf("Failed to run `CONFIG GET maxclients` command. %s", r.Err)
		return r.Err
	}

	res, err := r.Hash()
	if err != nil {
		logger.Errorf("Failed to fetch maxclients. %s", err)
		return err
	}

	maxsize, err := strconv.ParseFloat(res["maxclients"], 64)
	if err != nil {
		logger.Errorf("Failed to parse maxclients. %s", err)
		return err
	}

	stat["percentage_of_clients"] = 100.0 * stat["connected_clients"].(float64) / maxsize

	return nil
}

func calculateCapacity(c *redis.Client, stat map[string]interface{}) error {
	if err := fetchPercentageOfMemory(c, stat); err != nil {
		return err
	}
	if err := fetchPercentageOfClients(c, stat); err != nil {
		return err
	}
	return nil
}

// MetricKeyPrefix interface for PluginWithPrefix
func (m RedisPlugin) MetricKeyPrefix() string {
	if m.Prefix == "" {
		m.Prefix = "redis"
	}
	return m.Prefix
}

// FetchMetrics interface for mackerelplugin
func (m RedisPlugin) FetchMetrics() (map[string]interface{}, error) {
	network := "tcp"
	target := fmt.Sprintf("%s:%s", m.Host, m.Port)
	if m.Socket != "" {
		target = m.Socket
		network = "unix"
	}
	c, err := redis.DialTimeout(network, target, time.Duration(m.Timeout)*time.Second)
	if err != nil {
		logger.Errorf("Failed to connect redis. %s", err)
		return nil, err
	}
	defer c.Close()

	if m.Password != "" {
		if err = authenticateByPassword(c, m.Password); err != nil {
			return nil, err
		}
	}

	r := c.Cmd("info")
	if r.Err != nil {
		logger.Errorf("Failed to run info command. %s", r.Err)
		return nil, r.Err
	}
	str, err := r.Str()
	if err != nil {
		logger.Errorf("Failed to fetch information. %s", err)
		return nil, err
	}

	stat := make(map[string]interface{})

	keysStat := 0.0
	expiresStat := 0.0

	for _, line := range strings.Split(str, "\r\n") {
		if line == "" {
			continue
		}
		if re, _ := regexp.MatchString("^#", line); re {
			continue
		}

		record := strings.SplitN(line, ":", 2)
		if len(record) < 2 {
			continue
		}
		key, value := record[0], record[1]

		if re, _ := regexp.MatchString("^db", key); re {
			kv := strings.SplitN(value, ",", 3)
			keys, expires := kv[0], kv[1]

			keysKv := strings.SplitN(keys, "=", 2)
			keysFv, err := strconv.ParseFloat(keysKv[1], 64)
			if err != nil {
				logger.Warningf("Failed to parse db keys. %s", err)
			}
			keysStat += keysFv

			expiresKv := strings.SplitN(expires, "=", 2)
			expiresFv, err := strconv.ParseFloat(expiresKv[1], 64)
			if err != nil {
				logger.Warningf("Failed to parse db expires. %s", err)
			}
			expiresStat += expiresFv

			continue
		}

		stat[key], err = strconv.ParseFloat(value, 64)
		if err != nil {
			continue
		}
	}

	stat["keys"] = keysStat
	stat["expires"] = expiresStat

	if _, ok := stat["keys"]; !ok {
		stat["keys"] = 0
	}
	if _, ok := stat["expires"]; !ok {
		stat["expires"] = 0
	}

	if _, ok := stat["expired_keys"]; ok {
		stat["expired"] = stat["expired_keys"]
	} else {
		stat["expired"] = 0.0
	}

	if err := calculateCapacity(c, stat); err != nil {
		logger.Infof("Failed to calculate capacity. (The cause may be that AWS Elasticache Redis has no `CONFIG` command.) Skip these metrics. %s", err)
	}

	return stat, nil
}

// GraphDefinition interface for mackerelplugin
func (m RedisPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(m.Prefix)

	var graphdef = map[string]mp.Graphs{
		"queries": {
			Label: (labelPrefix + " Queries"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "total_commands_processed", Label: "Queries", Diff: true},
			},
		},
		"connections": {
			Label: (labelPrefix + " Connections"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "total_connections_received", Label: "Connections", Diff: true, Stacked: true},
				{Name: "rejected_connections", Label: "Rejected Connections", Diff: true, Stacked: true},
			},
		},
		"clients": {
			Label: (labelPrefix + " Clients"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "connected_clients", Label: "Connected Clients", Diff: false, Stacked: true},
				{Name: "blocked_clients", Label: "Blocked Clients", Diff: false, Stacked: true},
				{Name: "connected_slaves", Label: "Connected Slaves", Diff: false, Stacked: true},
			},
		},
		"keys": {
			Label: (labelPrefix + " Keys"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "keys", Label: "Keys", Diff: false},
				{Name: "expires", Label: "Keys with expiration", Diff: false},
				{Name: "expired", Label: "Expired Keys", Diff: false},
			},
		},
		"keyspace": {
			Label: (labelPrefix + " Keyspace"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "keyspace_hits", Label: "Keyspace Hits", Diff: true},
				{Name: "keyspace_misses", Label: "Keyspace Missed", Diff: true},
			},
		},
		"memory": {
			Label: (labelPrefix + " Memory"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "used_memory", Label: "Used Memory", Diff: false},
				{Name: "used_memory_rss", Label: "Used Memory RSS", Diff: false},
				{Name: "used_memory_peak", Label: "Used Memory Peak", Diff: false},
				{Name: "used_memory_lua", Label: "Used Memory Lua engine", Diff: false},
			},
		},
		"capacity": {
			Label: (labelPrefix + " Capacity"),
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				{Name: "percentage_of_memory", Label: "Percentage of memory", Diff: false},
				{Name: "percentage_of_clients", Label: "Percentage of clients", Diff: false},
			},
		},
	}

	return graphdef
}

// Do the plugin
func Do() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "6379", "Port")
	optPassowrd := flag.String("password", "", "Password")
	optSocket := flag.String("socket", "", "Server socket (overrides host and port)")
	optPrefix := flag.String("metric-key-prefix", "redis", "Metric key prefix")
	optTimeout := flag.Int("timeout", 5, "Timeout")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	redis := RedisPlugin{
		Timeout: *optTimeout,
		Prefix:  *optPrefix,
	}
	if *optSocket != "" {
		redis.Socket = *optSocket
	} else {
		redis.Host = *optHost
		redis.Port = *optPort
		redis.Password = *optPassowrd
	}
	helper := mp.NewMackerelPlugin(redis)
	helper.Tempfile = *optTempfile

	helper.Run()
}
<commit_msg>Fix redundant error in redis.go<commit_after>package mpredis

import (
	"flag"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/fzzy/radix/redis"
	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.redis")

// RedisPlugin mackerel plugin for Redis
type RedisPlugin struct {
	Host     string
	Port     string
	Password string
	Socket   string
	Prefix   string
	Timeout  int
	Tempfile string
}

func authenticateByPassword(c *redis.Client, password string) error {
	if r := c.Cmd("AUTH", password); r.Err != nil {
		logger.Errorf("Faild to authenticate. %s", r.Err)
		return r.Err
	}
	return nil
}

func fetchPercentageOfMemory(c *redis.Client, stat map[string]interface{}) error {
	r := c.Cmd("CONFIG", "GET", "maxmemory")
	if r.Err != nil {
		logger.Errorf("Failed to run `CONFIG GET maxmemory` command. %s", r.Err)
		return r.Err
	}

	res, err := r.Hash()
	if err != nil {
		logger.Errorf("Failed to fetch maxmemory. %s", err)
		return err
	}

	maxsize, err := strconv.ParseFloat(res["maxmemory"], 64)
	if err != nil {
		logger.Errorf("Failed to parse maxmemory. %s", err)
		return err
	}

	if maxsize == 0.0 {
		stat["percentage_of_memory"] = 0.0
	} else {
		stat["percentage_of_memory"] = 100.0 * stat["used_memory"].(float64) / maxsize
	}

	return nil
}

func fetchPercentageOfClients(c *redis.Client, stat map[string]interface{}) error {
	r := c.Cmd("CONFIG", "GET", "maxclients")
	if r.Err != nil {
		logger.Errorf("Failed to run `CONFIG GET maxclients` command. %s", r.Err)
		return r.Err
	}

	res, err := r.Hash()
	if err != nil {
		logger.Errorf("Failed to fetch maxclients. %s", err)
		return err
	}

	maxsize, err := strconv.ParseFloat(res["maxclients"], 64)
	if err != nil {
		logger.Errorf("Failed to parse maxclients. %s", err)
		return err
	}

	stat["percentage_of_clients"] = 100.0 * stat["connected_clients"].(float64) / maxsize

	return nil
}

func calculateCapacity(c *redis.Client, stat map[string]interface{}) error {
	if err := fetchPercentageOfMemory(c, stat); err != nil {
		return err
	}
	return fetchPercentageOfClients(c, stat)
}

// MetricKeyPrefix interface for PluginWithPrefix
func (m RedisPlugin) MetricKeyPrefix() string {
	if m.Prefix == "" {
		m.Prefix = "redis"
	}
	return m.Prefix
}

// FetchMetrics interface for mackerelplugin
func (m RedisPlugin) FetchMetrics() (map[string]interface{}, error) {
	network := "tcp"
	target := fmt.Sprintf("%s:%s", m.Host, m.Port)
	if m.Socket != "" {
		target = m.Socket
		network = "unix"
	}
	c, err := redis.DialTimeout(network, target, time.Duration(m.Timeout)*time.Second)
	if err != nil {
		logger.Errorf("Failed to connect redis. %s", err)
		return nil, err
	}
	defer c.Close()

	if m.Password != "" {
		if err = authenticateByPassword(c, m.Password); err != nil {
			return nil, err
		}
	}

	r := c.Cmd("info")
	if r.Err != nil {
		logger.Errorf("Failed to run info command. %s", r.Err)
		return nil, r.Err
	}
	str, err := r.Str()
	if err != nil {
		logger.Errorf("Failed to fetch information. %s", err)
		return nil, err
	}

	stat := make(map[string]interface{})

	keysStat := 0.0
	expiresStat := 0.0

	for _, line := range strings.Split(str, "\r\n") {
		if line == "" {
			continue
		}
		if re, _ := regexp.MatchString("^#", line); re {
			continue
		}

		record := strings.SplitN(line, ":", 2)
		if len(record) < 2 {
			continue
		}
		key, value := record[0], record[1]

		if re, _ := regexp.MatchString("^db", key); re {
			kv := strings.SplitN(value, ",", 3)
			keys, expires := kv[0], kv[1]

			keysKv := strings.SplitN(keys, "=", 2)
			keysFv, err := strconv.ParseFloat(keysKv[1], 64)
			if err != nil {
				logger.Warningf("Failed to parse db keys. %s", err)
			}
			keysStat += keysFv

			expiresKv := strings.SplitN(expires, "=", 2)
			expiresFv, err := strconv.ParseFloat(expiresKv[1], 64)
			if err != nil {
				logger.Warningf("Failed to parse db expires. %s", err)
			}
			expiresStat += expiresFv

			continue
		}

		stat[key], err = strconv.ParseFloat(value, 64)
		if err != nil {
			continue
		}
	}

	stat["keys"] = keysStat
	stat["expires"] = expiresStat

	if _, ok := stat["keys"]; !ok {
		stat["keys"] = 0
	}
	if _, ok := stat["expires"]; !ok {
		stat["expires"] = 0
	}

	if _, ok := stat["expired_keys"]; ok {
		stat["expired"] = stat["expired_keys"]
	} else {
		stat["expired"] = 0.0
	}

	if err := calculateCapacity(c, stat); err != nil {
		logger.Infof("Failed to calculate capacity. (The cause may be that AWS Elasticache Redis has no `CONFIG` command.) Skip these metrics. %s", err)
	}

	return stat, nil
}

// GraphDefinition interface for mackerelplugin
func (m RedisPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(m.Prefix)

	var graphdef = map[string]mp.Graphs{
		"queries": {
			Label: (labelPrefix + " Queries"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "total_commands_processed", Label: "Queries", Diff: true},
			},
		},
		"connections": {
			Label: (labelPrefix + " Connections"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "total_connections_received", Label: "Connections", Diff: true, Stacked: true},
				{Name: "rejected_connections", Label: "Rejected Connections", Diff: true, Stacked: true},
			},
		},
		"clients": {
			Label: (labelPrefix + " Clients"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "connected_clients", Label: "Connected Clients", Diff: false, Stacked: true},
				{Name: "blocked_clients", Label: "Blocked Clients", Diff: false, Stacked: true},
				{Name: "connected_slaves", Label: "Connected Slaves", Diff: false, Stacked: true},
			},
		},
		"keys": {
			Label: (labelPrefix + " Keys"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "keys", Label: "Keys", Diff: false},
				{Name: "expires", Label: "Keys with expiration", Diff: false},
				{Name: "expired", Label: "Expired Keys", Diff: false},
			},
		},
		"keyspace": {
			Label: (labelPrefix + " Keyspace"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "keyspace_hits", Label: "Keyspace Hits", Diff: true},
				{Name: "keyspace_misses", Label: "Keyspace Missed", Diff: true},
			},
		},
		"memory": {
			Label: (labelPrefix + " Memory"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "used_memory", Label: "Used Memory", Diff: false},
				{Name: "used_memory_rss", Label: "Used Memory RSS", Diff: false},
				{Name: "used_memory_peak", Label: "Used Memory Peak", Diff: false},
				{Name: "used_memory_lua", Label: "Used Memory Lua engine", Diff: false},
			},
		},
		"capacity": {
			Label: (labelPrefix + " Capacity"),
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				{Name: "percentage_of_memory", Label: "Percentage of memory", Diff: false},
				{Name: "percentage_of_clients", Label: "Percentage of clients", Diff: false},
			},
		},
	}

	return graphdef
}

// Do the plugin
func Do() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "6379", "Port")
	optPassowrd := flag.String("password", "", "Password")
	optSocket := flag.String("socket", "", "Server socket (overrides host and port)")
	optPrefix := flag.String("metric-key-prefix", "redis", "Metric key prefix")
	optTimeout := flag.Int("timeout", 5, "Timeout")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	redis := RedisPlugin{
		Timeout: *optTimeout,
		Prefix:  *optPrefix,
	}
	if *optSocket != "" {
		redis.Socket = *optSocket
	} else {
		redis.Host = *optHost
		redis.Port = *optPort
		redis.Password = *optPassowrd
	}
	helper := mp.NewMackerelPlugin(redis)
	helper.Tempfile = *optTempfile

	helper.Run()
}
<|endoftext|>


<commit_before>// Copyright 2013 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.

package main

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	yaml "gopkg.in/yaml.v2"
)

var (
	identifierRE   = `[a-zA-Z_][a-zA-Z0-9_]+`
	statsdMetricRE = `[a-zA-Z_](-?[a-zA-Z0-9_])+`

	metricLineRE = regexp.MustCompile(`^(\*\.|` + statsdMetricRE + `\.)+(\*|` + statsdMetricRE + `)$`)
	labelLineRE  = regexp.MustCompile(`^(` + identifierRE + `)\s*=\s*"(.*)"$`)
	metricNameRE = regexp.MustCompile(`^` + identifierRE + `$`)
)

type mapperConfigDefaults struct {
	TimerType timerType `yaml:"timer_type"`
	Buckets   []float64 `yaml:"buckets"`
}

type metricMapper struct {
	Defaults mapperConfigDefaults `yaml:"defaults"`
	Mappings []metricMapping      `yaml:"mappings"`
	mutex    sync.Mutex
}

type metricMapping struct {
	Match     string `yaml:"match"`
	regex     *regexp.Regexp
	Labels    prometheus.Labels `yaml:"labels"`
	TimerType timerType         `yaml:"timer_type"`
	Buckets   []float64         `yaml:"buckets"`
}

type configLoadStates int

const (
	SEARCHING configLoadStates = iota
	METRIC_DEFINITION
)

func (m *metricMapper) initFromString(fileContents string) error {
	lines := strings.Split(fileContents, "\n")
	numLines := len(lines)
	state := SEARCHING

	parsedMappings := []metricMapping{}
	currentMapping := metricMapping{Labels: prometheus.Labels{}}
	for i, line := range lines {
		line = strings.TrimSpace(line)

		// skip comments
		if strings.HasPrefix(line, "#") {
			continue
		}

		switch state {
		case SEARCHING:
			if line == "" {
				continue
			}
			if !metricLineRE.MatchString(line) {
				return fmt.Errorf("line %d: expected metric match line, got: %s", i, line)
			}

			// Translate the glob-style metric match line into a proper regex that we
			// can use to match metrics later on.
			metricRe := strings.Replace(line, ".", "\\.", -1)
			metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
			currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

			state = METRIC_DEFINITION

		case METRIC_DEFINITION:
			if (i == numLines-1) && (line != "") {
				return fmt.Errorf("Line %d: missing terminating newline", i)
			}
			if line == "" {
				if len(currentMapping.Labels) == 0 {
					return fmt.Errorf("Line %d: metric mapping didn't set any labels", i)
				}
				if _, ok := currentMapping.Labels["name"]; !ok {
					return fmt.Errorf("Line %d: metric mapping didn't set a metric name", i)
				}
				parsedMappings = append(parsedMappings, currentMapping)
				state = SEARCHING
				currentMapping = metricMapping{Labels: prometheus.Labels{}}
				continue
			}
			if err := m.updateMapping(line, i, &currentMapping); err != nil {
				return err
			}
		default:
			panic("illegal state")
		}
	}

	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.Mappings = parsedMappings

	mappingsCount.Set(float64(len(parsedMappings)))

	return nil
}

func (m *metricMapper) initFromYAMLString(fileContents string) error {
	var n metricMapper

	if err := yaml.Unmarshal([]byte(fileContents), &n); err != nil {
		return err
	}

	if n.Defaults.Buckets == nil || len(n.Defaults.Buckets) == 0 {
		n.Defaults.Buckets = prometheus.DefBuckets
	}

	for i := range n.Mappings {
		currentMapping := &n.Mappings[i]

		if !metricLineRE.MatchString(currentMapping.Match) {
			return fmt.Errorf("invalid match: %s", currentMapping.Match)
		}

		// Check that label is correct.
		for k, v := range currentMapping.Labels {
			label := fmt.Sprintf("%s=%q", k, v)
			if len(labelLineRE.FindStringSubmatch(label)) != 3 {
				return fmt.Errorf("invalid label: %s: %s", k, v)
			}
			if k == "name" && !metricNameRE.MatchString(v) {
				return fmt.Errorf("metric name '%s' doesn't match regex '%s'", v, metricNameRE)
			}
		}

		if _, ok := currentMapping.Labels["name"]; !ok {
			return fmt.Errorf("line %d: metric mapping didn't set a metric name", i)
		}

		// Translate the glob-style metric match line into a proper regex that we
		// can use to match metrics later on.
		metricRe := strings.Replace(currentMapping.Match, ".", "\\.", -1)
		metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
		currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

		if currentMapping.TimerType == "" {
			currentMapping.TimerType = n.Defaults.TimerType
		}

		if currentMapping.Buckets == nil || len(currentMapping.Buckets) == 0 {
			currentMapping.Buckets = n.Defaults.Buckets
		}

	}

	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.Defaults = n.Defaults
	m.Mappings = n.Mappings

	mappingsCount.Set(float64(len(n.Mappings)))

	return nil
}

func (m *metricMapper) initFromFile(fileName string) error {
	mappingStr, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}
	switch strings.ToLower(filepath.Ext(fileName)) {
	case ".yaml", ".yml":
		return m.initFromYAMLString(string(mappingStr))
	default:
		return m.initFromString(string(mappingStr))
	}
}

func (m *metricMapper) getMapping(statsdMetric string) (*metricMapping, prometheus.Labels, bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	for _, mapping := range m.Mappings {
		matches := mapping.regex.FindStringSubmatchIndex(statsdMetric)
		if len(matches) == 0 {
			continue
		}

		labels := prometheus.Labels{}
		for label, valueExpr := range mapping.Labels {
			value := mapping.regex.ExpandString([]byte{}, valueExpr, statsdMetric, matches)
			labels[label] = string(value)
		}
		return &mapping, labels, true
	}

	return nil, nil, false
}

func (m *metricMapper) updateMapping(line string, i int, mapping *metricMapping) error {
	matches := labelLineRE.FindStringSubmatch(line)
	if len(matches) != 3 {
		return fmt.Errorf("Line %d: expected label mapping line, got: %s", i, line)
	}
	label, value := matches[1], matches[2]
	if label == "name" && !metricNameRE.MatchString(value) {
		return fmt.Errorf("Line %d: metric name '%s' doesn't match regex '%s'", i, value, metricNameRE)
	}
	(*mapping).Labels[label] = value
	return nil
}
<commit_msg>simplify metric label key validation<commit_after>// Copyright 2013 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.

package main

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	yaml "gopkg.in/yaml.v2"
)

var (
	identifierRE   = `[a-zA-Z_][a-zA-Z0-9_]+`
	statsdMetricRE = `[a-zA-Z_](-?[a-zA-Z0-9_])+`

	metricLineRE = regexp.MustCompile(`^(\*\.|` + statsdMetricRE + `\.)+(\*|` + statsdMetricRE + `)$`)
	labelLineRE  = regexp.MustCompile(`^(` + identifierRE + `)\s*=\s*"(.*)"$`)
	metricNameRE = regexp.MustCompile(`^` + identifierRE + `$`)
)

type mapperConfigDefaults struct {
	TimerType timerType `yaml:"timer_type"`
	Buckets   []float64 `yaml:"buckets"`
}

type metricMapper struct {
	Defaults mapperConfigDefaults `yaml:"defaults"`
	Mappings []metricMapping      `yaml:"mappings"`
	mutex    sync.Mutex
}

type metricMapping struct {
	Match     string `yaml:"match"`
	regex     *regexp.Regexp
	Labels    prometheus.Labels `yaml:"labels"`
	TimerType timerType         `yaml:"timer_type"`
	Buckets   []float64         `yaml:"buckets"`
}

type configLoadStates int

const (
	SEARCHING configLoadStates = iota
	METRIC_DEFINITION
)

func (m *metricMapper) initFromString(fileContents string) error {
	lines := strings.Split(fileContents, "\n")
	numLines := len(lines)
	state := SEARCHING

	parsedMappings := []metricMapping{}
	currentMapping := metricMapping{Labels: prometheus.Labels{}}
	for i, line := range lines {
		line = strings.TrimSpace(line)

		// skip comments
		if strings.HasPrefix(line, "#") {
			continue
		}

		switch state {
		case SEARCHING:
			if line == "" {
				continue
			}
			if !metricLineRE.MatchString(line) {
				return fmt.Errorf("line %d: expected metric match line, got: %s", i, line)
			}

			// Translate the glob-style metric match line into a proper regex that we
			// can use to match metrics later on.
			metricRe := strings.Replace(line, ".", "\\.", -1)
			metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
			currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

			state = METRIC_DEFINITION

		case METRIC_DEFINITION:
			if (i == numLines-1) && (line != "") {
				return fmt.Errorf("Line %d: missing terminating newline", i)
			}
			if line == "" {
				if len(currentMapping.Labels) == 0 {
					return fmt.Errorf("Line %d: metric mapping didn't set any labels", i)
				}
				if _, ok := currentMapping.Labels["name"]; !ok {
					return fmt.Errorf("Line %d: metric mapping didn't set a metric name", i)
				}
				parsedMappings = append(parsedMappings, currentMapping)
				state = SEARCHING
				currentMapping = metricMapping{Labels: prometheus.Labels{}}
				continue
			}
			if err := m.updateMapping(line, i, &currentMapping); err != nil {
				return err
			}
		default:
			panic("illegal state")
		}
	}

	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.Mappings = parsedMappings

	mappingsCount.Set(float64(len(parsedMappings)))

	return nil
}

func (m *metricMapper) initFromYAMLString(fileContents string) error {
	var n metricMapper

	if err := yaml.Unmarshal([]byte(fileContents), &n); err != nil {
		return err
	}

	if n.Defaults.Buckets == nil || len(n.Defaults.Buckets) == 0 {
		n.Defaults.Buckets = prometheus.DefBuckets
	}

	for i := range n.Mappings {
		currentMapping := &n.Mappings[i]

		if !metricLineRE.MatchString(currentMapping.Match) {
			return fmt.Errorf("invalid match: %s", currentMapping.Match)
		}

		// Check that label is correct.
		for k, v := range currentMapping.Labels {
			if !metricNameRE.MatchString(k) {
				return fmt.Errorf("invalid label key: %s", k)
			}
			if k == "name" && !metricNameRE.MatchString(v) {
				return fmt.Errorf("metric name '%s' doesn't match regex '%s'", v, metricNameRE)
			}
		}

		if _, ok := currentMapping.Labels["name"]; !ok {
			return fmt.Errorf("line %d: metric mapping didn't set a metric name", i)
		}

		// Translate the glob-style metric match line into a proper regex that we
		// can use to match metrics later on.
		metricRe := strings.Replace(currentMapping.Match, ".", "\\.", -1)
		metricRe = strings.Replace(metricRe, "*", "([^.]+)", -1)
		currentMapping.regex = regexp.MustCompile("^" + metricRe + "$")

		if currentMapping.TimerType == "" {
			currentMapping.TimerType = n.Defaults.TimerType
		}

		if currentMapping.Buckets == nil || len(currentMapping.Buckets) == 0 {
			currentMapping.Buckets = n.Defaults.Buckets
		}

	}

	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.Defaults = n.Defaults
	m.Mappings = n.Mappings

	mappingsCount.Set(float64(len(n.Mappings)))

	return nil
}

func (m *metricMapper) initFromFile(fileName string) error {
	mappingStr, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}
	switch strings.ToLower(filepath.Ext(fileName)) {
	case ".yaml", ".yml":
		return m.initFromYAMLString(string(mappingStr))
	default:
		return m.initFromString(string(mappingStr))
	}
}

func (m *metricMapper) getMapping(statsdMetric string) (*metricMapping, prometheus.Labels, bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	for _, mapping := range m.Mappings {
		matches := mapping.regex.FindStringSubmatchIndex(statsdMetric)
		if len(matches) == 0 {
			continue
		}

		labels := prometheus.Labels{}
		for label, valueExpr := range mapping.Labels {
			value := mapping.regex.ExpandString([]byte{}, valueExpr, statsdMetric, matches)
			labels[label] = string(value)
		}
		return &mapping, labels, true
	}

	return nil, nil, false
}

func (m *metricMapper) updateMapping(line string, i int, mapping *metricMapping) error {
	matches := labelLineRE.FindStringSubmatch(line)
	if len(matches) != 3 {
		return fmt.Errorf("Line %d: expected label mapping line, got: %s", i, line)
	}
	label, value := matches[1], matches[2]
	if label == "name" && !metricNameRE.MatchString(value) {
		return fmt.Errorf("Line %d: metric name '%s' doesn't match regex '%s'", i, value, metricNameRE)
	}
	(*mapping).Labels[label] = value
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"html/template"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"

	"golang.org/x/crypto/ssh/terminal"
	yaml "gopkg.in/yaml.v2"
)

type marvin struct {
	Config struct {
		Relative  bool   `yaml:"relative"`
		Delemiter string `yaml:"delemiter"`
	}
	Tasks     map[string]string `yaml:"tasks"`
	Inventory struct {
		static    []map[string]string
		StaticRaw string            `yaml:"static"`
		Dynamic   map[string]string `yaml:"dynamic"`
	} `yaml:"inventory"`
	filtered []map[string]string
	lock     *sync.Mutex
}

func newMarvin(config, currentDirectory, query, task, args string) *marvin {
	m := &marvin{
		lock: &sync.Mutex{},
	}

	// set configuration defaults
	m.setConfigDefaults()

	// set task defaults
	m.setTaskDefaults()

	// set some default dynamic inventory
	m.setDynamicInventoryDefaults()

	err := yaml.Unmarshal([]byte(config), &m)
	if err != nil {
		speak(err.Error(), false)
		speak("marvin.yml> invalid", true)
	}

	if m.Config.Relative {
		os.Chdir(currentDirectory)
	}

	// stdin to static inventory
	if !terminal.IsTerminal(0) {
		stdin, _ := ioutil.ReadAll(os.Stdin)
		m.Inventory.StaticRaw += "\n" + string(stdin)
	}

	m.filter(query)

	if len(m.filtered) == 0 {
		speak("inventory> none matched your criteria", true)
	}
	m.task(task, args)
	return m
}

func (m *marvin) rawToInventory(raw string) []map[string]string {
	inventory := make([]map[string]string, 0)
	// convert strings to actual inventory
	for _, iRecord := range strings.Split(strings.TrimSpace(m.Inventory.StaticRaw), "\n") {
		for _, row := range strings.Split(iRecord, "\n") {
			i := make(map[string]string)
			if row == "" {
				continue
			}
			for id, kvs := range strings.Split(row, m.Config.Delemiter) {
				kv := strings.SplitN(kvs, ":", 2)
				if len(kv) == 1 {
					i["id"] = kv[0]
				} else {
					i[kv[0]] = strings.Join(kv[1:], m.Config.Delemiter)
					// always set an id
					if id == 0 {
						i["id"] = i[kv[0]]
					}
				}

				// regardless, lets add a raw
				i["raw"] = row
			}

			// add the inventory
			inventory = append(inventory, i)
		}
	}
	m.Inventory.static = inventory
	return m.Inventory.static
}

func (m *marvin) filter(queryString string) []map[string]string {
	var sKey, sValue string
	// kiss for now, replace * with .* and then regex match
	filtered := make([]map[string]string, 0)
	// figure out what we are searching for.
	q := strings.Split(queryString, ":")
	if len(q) == 2 {
		sKey, sValue = q[0], q[1]
		if sKey == "" {
			sKey = "id"
		}
		if sValue == "" {
			sValue = "*"
		}
	} else {
		sKey, sValue = "id", q[0]
	}

	if dynamicCmd, dynamicCmdExists := m.Inventory.Dynamic[sKey]; dynamicCmdExists {
		dInventoryOutput, dInventoryError := m.exec(dynamicCmd)
		if dInventoryError == nil && strings.TrimSpace(dInventoryOutput) != "" {
			// if no errors, add it
			m.lock.Lock()
			m.Inventory.StaticRaw += "\n"
			m.lock.Unlock()
			for _, row := range strings.Split(dInventoryOutput, "\n") {
				// add each record, with the inventory name
				m.lock.Lock()
				m.Inventory.StaticRaw += sKey + ":" + row + "\n"
				m.lock.Unlock()
			}
		}
	}

	var wg sync.WaitGroup
	if sKey == "id" || sKey == "raw" {
		// sigh, we need to generate ALL inventory because we just don't know
		for dInventoryName, dInventoryCmd := range m.Inventory.Dynamic {
			go func(dInventoryCmd, dInventoryName string) {
				wg.Add(1)
				m.lock.Lock()
				m.Inventory.StaticRaw += "\n"
				m.lock.Unlock()
				if dInventoryOutput, dInventoryError := m.exec(dInventoryCmd); dInventoryError == nil && strings.TrimSpace(dInventoryOutput) != "" {
					// no errors, sweet
					for _, row := range strings.Split(dInventoryOutput, "\n") {
						if row == "" {
							continue
						}
						// add each record, with the inventory name
						m.lock.Lock()
						m.Inventory.StaticRaw += dInventoryName + ":" + row + "\n"
						m.lock.Unlock()
					}
				}
				wg.Done()
			}(dInventoryCmd, dInventoryName)
		}
		wg.Wait()
	}

	// inventory should now be set ...
	m.rawToInventory(m.Inventory.StaticRaw)

	// replace
	sKey = strings.Replace(sKey, "*", ".*", -1)
	sValue = strings.Replace(sValue, "*", ".*", -1)

	for _, i := range m.Inventory.static {
		for k, v := range i {
			keyMatch, _ := regexp.MatchString(sKey, k)
			valueMatch, _ := regexp.MatchString(sValue, v)
			if keyMatch && valueMatch {
				filtered = append(filtered, i)
			}
		}
	}
	m.filtered = filtered
	return filtered
}

func (m *marvin) task(taskName, args string) {
	rawTask, exists := m.Tasks[taskName]
	var wg sync.WaitGroup
	if !exists {
		rawTask = "sh -c '" + taskName + " " + args + "'"
	}

	// first lets do an inventory check
	rawCommands := m.inventoryCheck(rawTask, m.filtered, args)

	// sweet, moving along
	// TODO: we are just going to spin them all up(we will figure out batching later)
	for cmdID, command := range rawCommands {
		wg.Add(1)
		go func(cmdID, command string) {

			output, execError := m.exec(command)

			// hopefully only one line :fingers_crossed
			if len(strings.Split(output, "\n")) > 1 {
				output = "\n" + output
			}

			if execError != nil {
				color.Red(cmdID + "> " + output)
			} else {
				color.Green(cmdID + "> " + output)
			}
			wg.Done()
		}(cmdID, command)
	}
	wg.Wait()
}

func (m *marvin) inventoryCheck(task string, filtered []map[string]string, args string) map[string]string {
	task = strings.Replace(task, "<", "<", -1)
	rawCommands := make(map[string]string, 0)
	for _, inventory := range filtered {
		// add defaults to filtered inventory
		inventory["args"] = args
		inventory["time"] = time.Now().String()
		cmd, cmdParseError := m.template(task, inventory)
		if cmdParseError != nil {
			speak(inventory["id"]+"> missing/invalid arguments", true)
		}
		rawCommands[inventory["id"]] = cmd
	}
	return rawCommands
}

func (m *marvin) template(task string, inventory map[string]string) (string, error) {
	template := template.Must(template.New("translate").Parse(task))
	b := new(bytes.Buffer)
	err := template.Execute(b, inventory)
	if err == nil {
		return b.String(), nil
	}
	return "", fmt.Errorf("id: " + inventory["id"])
}

func (m *marvin) exec(command string) (string, error) {
	cmd := exec.Command("sh", "-c", command)
	execOutput, execError := cmd.CombinedOutput()
	output := strings.TrimSpace(string(execOutput))
	output = strings.Trim(output, "\r\n")
	return output, execError
}

func (m *marvin) setDynamicInventoryDefaults() {
	m.Inventory.Dynamic = make(map[string]string)

	m.Inventory.Dynamic["file"] = "ls -R -1"
	m.Inventory.Dynamic["dir"] = "ls -d */ | cut -f1 -d'/'"
	m.Inventory.Dynamic["branch"] = "git branch | cut -c 3-"
	m.Inventory.Dynamic["docker"] = "docker ps --format \"{{ .Names }}\""
	m.Inventory.Dynamic["bookmark"] = `cat ~/Library/Application\ Support/Google/Chrome/Default/Bookmarks | grep -i http | sed "s/ //g" | sed "s/\"//g" | sed "s/url://g"`
	m.Inventory.Dynamic["ec2"] = `aws ec2 describe-instances --output=text | grep "TAGS$(printf '\t')Name" | grep -v " " | sed "s/TAGS$(printf '\t')Name$(printf '\t')//g"`
}

func (m *marvin) setTaskDefaults() {
	m.Tasks = make(map[string]string)
	// set default tasks
	m.Tasks["ssh"] = "ssh {{ .host }} {{ .args }}"
	m.Tasks["ls"] = "echo {{ .raw }}"
	m.Tasks["open"] = "open {{ .id }} && echo opening"
}

func (m *marvin) setConfigDefaults() {
	// delemiter
	m.Config.Delemiter = " "

	// relative?
	m.Config.Relative = true
}
<commit_msg>fixed a bug with the wait groups<commit_after>package main

import (
	"bytes"
	"fmt"
	"html/template"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"

	"golang.org/x/crypto/ssh/terminal"
	yaml "gopkg.in/yaml.v2"
)

type marvin struct {
	Config struct {
		Relative  bool   `yaml:"relative"`
		Delemiter string `yaml:"delemiter"`
	}
	Tasks     map[string]string `yaml:"tasks"`
	Inventory struct {
		static    []map[string]string
		StaticRaw string            `yaml:"static"`
		Dynamic   map[string]string `yaml:"dynamic"`
	} `yaml:"inventory"`
	filtered []map[string]string
	lock     *sync.Mutex
}

func newMarvin(config, currentDirectory, query, task, args string) *marvin {
	m := &marvin{
		lock: &sync.Mutex{},
	}

	// set configuration defaults
	m.setConfigDefaults()

	// set task defaults
	m.setTaskDefaults()

	// set some default dynamic inventory
	m.setDynamicInventoryDefaults()

	err := yaml.Unmarshal([]byte(config), &m)
	if err != nil {
		speak(err.Error(), false)
		speak("marvin.yml> invalid", true)
	}

	if m.Config.Relative {
		os.Chdir(currentDirectory)
	}

	// stdin to static inventory
	if !terminal.IsTerminal(0) {
		stdin, _ := ioutil.ReadAll(os.Stdin)
		m.Inventory.StaticRaw += "\n" + string(stdin)
	}

	m.filter(query)

	if len(m.filtered) == 0 {
		speak("inventory> none matched your criteria", true)
	}
	m.task(task, args)
	return m
}

func (m *marvin) rawToInventory(raw string) []map[string]string {
	inventory := make([]map[string]string, 0)
	// convert strings to actual inventory
	for _, iRecord := range strings.Split(strings.TrimSpace(m.Inventory.StaticRaw), "\n") {
		for _, row := range strings.Split(iRecord, "\n") {
			i := make(map[string]string)
			if row == "" {
				continue
			}
			for id, kvs := range strings.Split(row, m.Config.Delemiter) {
				kv := strings.SplitN(kvs, ":", 2)
				if len(kv) == 1 {
					i["id"] = kv[0]
				} else {
					i[kv[0]] = strings.Join(kv[1:], m.Config.Delemiter)
					// always set an id
					if id == 0 {
						i["id"] = i[kv[0]]
					}
				}

				// regardless, lets add a raw
				i["raw"] = row
			}

			// add the inventory
			inventory = append(inventory, i)
		}
	}
	m.Inventory.static = inventory
	return m.Inventory.static
}

func (m *marvin) filter(queryString string) []map[string]string {
	var sKey, sValue string
	// kiss for now, replace * with .* and then regex match
	filtered := make([]map[string]string, 0)
	// figure out what we are searching for.
	q := strings.Split(queryString, ":")
	if len(q) == 2 {
		sKey, sValue = q[0], q[1]
		if sKey == "" {
			sKey = "id"
		}
		if sValue == "" {
			sValue = "*"
		}
	} else {
		sKey, sValue = "id", q[0]
	}

	if dynamicCmd, dynamicCmdExists := m.Inventory.Dynamic[sKey]; dynamicCmdExists {
		dInventoryOutput, dInventoryError := m.exec(dynamicCmd)
		if dInventoryError == nil && strings.TrimSpace(dInventoryOutput) != "" {
			// if no errors, add it
			m.lock.Lock()
			m.Inventory.StaticRaw += "\n"
			m.lock.Unlock()
			for _, row := range strings.Split(dInventoryOutput, "\n") {
				// add each record, with the inventory name
				m.lock.Lock()
				m.Inventory.StaticRaw += sKey + ":" + row + "\n"
				m.lock.Unlock()
			}
		}
	}

	var wg sync.WaitGroup
	if sKey == "id" || sKey == "raw" {
		// sigh, we need to generate ALL inventory because we just don't know
		for dInventoryName, dInventoryCmd := range m.Inventory.Dynamic {
			wg.Add(1)
			go func(dInventoryCmd, dInventoryName string) {
				m.lock.Lock()
				m.Inventory.StaticRaw += "\n"
				m.lock.Unlock()
				if dInventoryOutput, dInventoryError := m.exec(dInventoryCmd); dInventoryError == nil && strings.TrimSpace(dInventoryOutput) != "" {
					// no errors, sweet
					for _, row := range strings.Split(dInventoryOutput, "\n") {
						if row == "" {
							continue
						}
						// add each record, with the inventory name
						m.lock.Lock()
						m.Inventory.StaticRaw += dInventoryName + ":" + row + "\n"
						m.lock.Unlock()
					}
				}
				wg.Done()
			}(dInventoryCmd, dInventoryName)
		}
		wg.Wait()
	}

	// inventory should now be set ...
	m.rawToInventory(m.Inventory.StaticRaw)

	// replace
	sKey = strings.Replace(sKey, "*", ".*", -1)
	sValue = strings.Replace(sValue, "*", ".*", -1)

	for _, i := range m.Inventory.static {
		for k, v := range i {
			keyMatch, _ := regexp.MatchString(sKey, k)
			valueMatch, _ := regexp.MatchString(sValue, v)
			if keyMatch && valueMatch {
				filtered = append(filtered, i)
			}
		}
	}
	m.filtered = filtered
	return filtered
}

func (m *marvin) task(taskName, args string) {
	rawTask, exists := m.Tasks[taskName]
	var wg sync.WaitGroup
	if !exists {
		rawTask = "sh -c '" + taskName + " " + args + "'"
	}

	// first lets do an inventory check
	rawCommands := m.inventoryCheck(rawTask, m.filtered, args)

	// sweet, moving along
	// TODO: we are just going to spin them all up(we will figure out batching later)
	for cmdID, command := range rawCommands {
		wg.Add(1)
		go func(cmdID, command string) {

			output, execError := m.exec(command)

			// hopefully only one line :fingers_crossed
			if len(strings.Split(output, "\n")) > 1 {
				output = "\n" + output
			}

			if execError != nil {
				color.Red(cmdID + "> " + output)
			} else {
				color.Green(cmdID + "> " + output)
			}
			wg.Done()
		}(cmdID, command)
	}
	wg.Wait()
}

func (m *marvin) inventoryCheck(task string, filtered []map[string]string, args string) map[string]string {
	task = strings.Replace(task, "<", "<", -1)
	rawCommands := make(map[string]string, 0)
	for _, inventory := range filtered {
		// add defaults to filtered inventory
		inventory["args"] = args
		inventory["time"] = time.Now().String()
		cmd, cmdParseError := m.template(task, inventory)
		if cmdParseError != nil {
			speak(inventory["id"]+"> missing/invalid arguments", true)
		}
		rawCommands[inventory["id"]] = cmd
	}
	return rawCommands
}

func (m *marvin) template(task string, inventory map[string]string) (string, error) {
	template := template.Must(template.New("translate").Parse(task))
	b := new(bytes.Buffer)
	err := template.Execute(b, inventory)
	if err == nil {
		return b.String(), nil
	}
	return "", fmt.Errorf("id: " + inventory["id"])
}

func (m *marvin) exec(command string) (string, error) {
	cmd := exec.Command("sh", "-c", command)
	execOutput, execError := cmd.CombinedOutput()
	output := strings.TrimSpace(string(execOutput))
	output = strings.Trim(output, "\r\n")
	return output, execError
}

func (m *marvin) setDynamicInventoryDefaults() {
	m.Inventory.Dynamic = make(map[string]string)

	m.Inventory.Dynamic["file"] = "ls -R -1"
	m.Inventory.Dynamic["dir"] = "ls -d */ | cut -f1 -d'/'"
	m.Inventory.Dynamic["branch"] = "git branch | cut -c 3-"
	m.Inventory.Dynamic["docker"] = "docker ps --format \"{{ .Names }}\""
	m.Inventory.Dynamic["bookmark"] = `cat ~/Library/Application\ Support/Google/Chrome/Default/Bookmarks | grep -i http | sed "s/ //g" | sed "s/\"//g" | sed "s/url://g"`
	m.Inventory.Dynamic["ec2"] = `aws ec2 describe-instances --output=text | grep "TAGS$(printf '\t')Name" | grep -v " " | sed "s/TAGS$(printf '\t')Name$(printf '\t')//g"`
}

func (m *marvin) setTaskDefaults() {
	m.Tasks = make(map[string]string)
	// set default tasks
	m.Tasks["ssh"] = "ssh {{ .host }} {{ .args }}"
	m.Tasks["ls"] = "echo {{ .raw }}"
	m.Tasks["open"] = "open {{ .id }} && echo opening"
}

func (m *marvin) setConfigDefaults() {
	// delemiter
	m.Config.Delemiter = " "

	// relative?
	m.Config.Relative = true
}
<|endoftext|>


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

package uid

import (
	"crypto/rand"
	"errors"
	"fmt"
	"log"
	"math"
	mrand "math/rand"
	"strings"
	"sync"
	"time"

	"context"

	"github.com/dgryski/go-farm"

	"github.com/dgraph-io/dgraph/posting"
	"github.com/dgraph-io/dgraph/posting/types"
	"github.com/dgraph-io/dgraph/store"
	"github.com/dgraph-io/dgraph/x"
)

var lmgr *lockManager
var letterRunes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ "
var uidStore *store.Store
var eidPool = sync.Pool{
	New: func() interface{} {
		return new(entry)
	},
}

type entry struct {
	sync.Mutex
	ts time.Time
}

func (e *entry) isOld() bool {
	e.Lock()
	defer e.Unlock()
	return time.Now().Sub(e.ts) > time.Minute
}

type lockManager struct {
	sync.RWMutex
	locks map[string]*entry
}

func (lm *lockManager) newOrExisting(xid string) *entry {
	lm.RLock()
	if e, ok := lm.locks[xid]; ok {
		lm.RUnlock()
		return e
	}
	lm.RUnlock()

	lm.Lock()
	defer lm.Unlock()
	if e, ok := lm.locks[xid]; ok {
		return e
	}
	e := eidPool.Get().(*entry)
	e.ts = time.Now()
	lm.locks[xid] = e
	return e
}

func (lm *lockManager) clean() {
	ticker := time.NewTicker(time.Minute)
	for _ = range ticker.C {
		count := 0
		lm.Lock()
		for xid, e := range lm.locks {
			if e.isOld() {
				count += 1
				delete(lm.locks, xid)
				eidPool.Put(e)
			}
		}
		lm.Unlock()
		// A minute is enough to avoid the race condition issue for
		// uid allocation to an xid.
	}
}

// package level init
func init() {
	lmgr = new(lockManager)
	lmgr.locks = make(map[string]*entry)
	// TODO(manishrjain): This map should be cleaned up.
	// go lmgr.clean()
}

func Init(ps *store.Store) {
	uidStore = ps
}

// allocateUniqueUid returns an integer in range:
// [minIdx, maxIdx] derived based on numInstances and instanceIdx.
// which hasn't already been allocated to other xids. It does this by
// taking the fingerprint of the xid appended with zero or more spaces
// until the obtained integer is unique.
func allocateUniqueUid(xid string, instanceIdx uint64,
	numInstances uint64) (uid uint64, rerr error) {
	mrand.Seed(time.Now().UnixNano())
	mod := math.MaxUint64 / numInstances
	minIdx := instanceIdx * mod
	txid := []byte(xid)
	val := xid
	if strings.HasPrefix(xid, "_new_:") {
		if _, err := rand.Read(txid); err != nil {
			return 0, err
		}

		val = "_new_"
	}

	for ; ; txid = append(txid, letterRunes[mrand.Intn(len(letterRunes))]) {
		uid1 := farm.Fingerprint64(txid) // Generate from hash.
		uid = (uid1 % mod) + minIdx
		if uid == math.MaxUint64 {
			continue
		}

		// Check if this uid has already been allocated.
		key := posting.Key(uid, "_xid_") // uid -> "_xid_" -> xid
		pl := posting.GetOrCreate(key, uidStore)

		if pl.Length() > 0 {
			// Something already present here.
			var p types.Posting
			pl.Get(&p, 0)
			continue
		}

		// Uid hasn't been assigned yet.
		t := x.DirectedEdge{
			Value:     []byte(val), // not txid
			Source:    "_assigner_",
			Timestamp: time.Now(),
		}
		_, rerr = pl.AddMutation(context.Background(), t, posting.Set)
		return uid, rerr
	}
}

// assignNew assigns a uid to a given xid and writes it to the
// posting list.
func assignNew(pl *posting.List, xid string, instanceIdx uint64,
	numInstances uint64) (uint64, error) {

	entry := lmgr.newOrExisting(xid)
	entry.Lock()
	entry.ts = time.Now()
	defer entry.Unlock()

	if pl.Length() > 1 {
		log.Fatalf("We shouldn't have more than 1 uid for xid: %v\n", xid)

	} else if pl.Length() > 0 {
		var p types.Posting
		if ok := pl.Get(&p, 0); !ok {
			return 0, errors.New("While retrieving entry from posting list.")
		}
		return p.Uid(), nil
	}

	// No current id exists. Create one.
	uid, err := allocateUniqueUid(xid, instanceIdx, numInstances)
	if err != nil {
		return 0, err
	}

	t := x.DirectedEdge{
		ValueId:   uid,
		Source:    "_assigner_",
		Timestamp: time.Now(),
	}
	_, rerr := pl.AddMutation(context.Background(), t, posting.Set)
	return uid, rerr
}

func StringKey(xid string) []byte {
	return []byte("_uid_|" + xid)
}

// Get returns the uid of the corresponding xid.
func Get(xid string) (uid uint64, rerr error) {
	key := StringKey(xid)
	pl := posting.GetOrCreate(key, uidStore)
	if pl.Length() == 0 {
		return 0, fmt.Errorf("xid: %v doesn't have any uid assigned.", xid)
	}
	if pl.Length() > 1 {
		log.Fatalf("We shouldn't have more than 1 uid for xid: %v\n", xid)
	}
	var p types.Posting
	if ok := pl.Get(&p, 0); !ok {
		return 0, fmt.Errorf("While retrieving entry from posting list")
	}
	return p.Uid(), nil
}

// GetOrAssign returns a unique integer (uid) for a given xid if
// it already exists or assigns a new uid and returns it.
func GetOrAssign(xid string, instanceIdx uint64,
	numInstances uint64) (uid uint64, rerr error) {

	// Prefix _new_ requires us to create a new uid(entity).
	if strings.HasPrefix(xid, "_new_:") {
		return allocateUniqueUid(xid, instanceIdx, numInstances)
	}

	key := StringKey(xid)
	pl := posting.GetOrCreate(key, uidStore)
	if pl.Length() == 0 {
		return assignNew(pl, xid, instanceIdx, numInstances)
	}

	if pl.Length() > 1 {
		log.Fatalf("We shouldn't have more than 1 uid for xid: %v\n", xid)
	} else {
		// We found one posting.
		var p types.Posting
		if ok := pl.Get(&p, 0); !ok {
			return 0, errors.New("While retrieving entry from posting list")
		}
		return p.Uid(), nil
	}
	return 0, errors.New("Some unhandled route lead me here." +
		" Wake the stupid developer up.")
}

// ExternalId returns the xid of a given uid by reading from the uidstore.
// It returns an error if there is no corresponding xid.
func ExternalId(uid uint64) (xid string, rerr error) {
	key := posting.Key(uid, "_xid_") // uid -> "_xid_" -> xid
	pl := posting.GetOrCreate(key, uidStore)
	if pl.Length() == 0 {
		return "", errors.New("NO external id")
	}

	if pl.Length() > 1 {
		log.Fatalf("This shouldn't be happening. Uid: %v", uid)
		return "", errors.New("Multiple external ids for this uid.")
	}

	var p types.Posting
	if ok := pl.Get(&p, 0); !ok {
		return "", errors.New("While retrieving posting")
	}

	if p.Uid() != math.MaxUint64 {
		log.Fatalf("Value uid must be MaxUint64. Uid: %v", p.Uid())
	}
	return string(p.ValueBytes()), rerr
}
<commit_msg>Use math/rand to aid with generating UIDs<commit_after>/*
 * Copyright 2015 DGraph Labs, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * 		http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package uid

import (
	"context"
	"errors"
	"fmt"
	"log"
	"math"
	"math/rand"
	"strings"
	"sync"
	"time"

	"github.com/dgryski/go-farm"

	"github.com/dgraph-io/dgraph/posting"
	"github.com/dgraph-io/dgraph/posting/types"
	"github.com/dgraph-io/dgraph/store"
	"github.com/dgraph-io/dgraph/x"
)

var lmgr *lockManager
var uidStore *store.Store
var eidPool = sync.Pool{
	New: func() interface{} {
		return new(entry)
	},
}

type entry struct {
	sync.Mutex
	ts time.Time
}

func (e *entry) isOld() bool {
	e.Lock()
	defer e.Unlock()
	return time.Now().Sub(e.ts) > time.Minute
}

type lockManager struct {
	sync.RWMutex
	locks map[string]*entry
}

func (lm *lockManager) newOrExisting(xid string) *entry {
	lm.RLock()
	if e, ok := lm.locks[xid]; ok {
		lm.RUnlock()
		return e
	}
	lm.RUnlock()

	lm.Lock()
	defer lm.Unlock()
	if e, ok := lm.locks[xid]; ok {
		return e
	}
	e := eidPool.Get().(*entry)
	e.ts = time.Now()
	lm.locks[xid] = e
	return e
}

func (lm *lockManager) clean() {
	ticker := time.NewTicker(time.Minute)
	for _ = range ticker.C {
		count := 0
		lm.Lock()
		for xid, e := range lm.locks {
			if e.isOld() {
				count += 1
				delete(lm.locks, xid)
				eidPool.Put(e)
			}
		}
		lm.Unlock()
		// A minute is enough to avoid the race condition issue for
		// uid allocation to an xid.
	}
}

// package level init
func init() {
	rand.Seed(time.Now().UnixNano())
	lmgr = new(lockManager)
	lmgr.locks = make(map[string]*entry)
	// TODO(manishrjain): This map should be cleaned up.
	// go lmgr.clean()
}

func Init(ps *store.Store) {
	uidStore = ps
}

// allocateUniqueUid returns an integer in range:
// [minIdx, maxIdx] derived based on numInstances and instanceIdx.
// which hasn't already been allocated to other xids. It does this by
// taking the fingerprint of the xid appended with zero or more spaces
// until the obtained integer is unique.
func allocateUniqueUid(xid string, instanceIdx uint64,
	numInstances uint64) (uid uint64, rerr error) {
	mod := math.MaxUint64 / numInstances
	minIdx := instanceIdx * mod

	txid := []byte(xid)
	val := xid
	if strings.HasPrefix(xid, "_new_:") {
		if _, err := rand.Read(txid); err != nil {
			return 0, err
		}
		val = "_new_"
	}

	suffix := make([]byte, 1)
	for i := 0; ; i++ {
		if i > 0 {
			_, err := rand.Read(suffix)
			x.Check(err)
			txid = append(txid, suffix[0])
		}

		uid1 := farm.Fingerprint64(txid) // Generate from hash.
		uid = (uid1 % mod) + minIdx
		if uid == math.MaxUint64 {
			continue
		}

		// Check if this uid has already been allocated.
		key := posting.Key(uid, "_xid_") // uid -> "_xid_" -> xid
		pl := posting.GetOrCreate(key, uidStore)

		if pl.Length() > 0 {
			// Something already present here.
			var p types.Posting
			pl.Get(&p, 0)
			continue
		}

		// Uid hasn't been assigned yet.
		t := x.DirectedEdge{
			Value:     []byte(val), // not txid
			Source:    "_assigner_",
			Timestamp: time.Now(),
		}
		_, rerr = pl.AddMutation(context.Background(), t, posting.Set)
		return uid, rerr
	}
}

// assignNew assigns a uid to a given xid and writes it to the
// posting list.
func assignNew(pl *posting.List, xid string, instanceIdx uint64,
	numInstances uint64) (uint64, error) {

	entry := lmgr.newOrExisting(xid)
	entry.Lock()
	entry.ts = time.Now()
	defer entry.Unlock()

	if pl.Length() > 1 {
		log.Fatalf("We shouldn't have more than 1 uid for xid: %v\n", xid)

	} else if pl.Length() > 0 {
		var p types.Posting
		if ok := pl.Get(&p, 0); !ok {
			return 0, errors.New("While retrieving entry from posting list.")
		}
		return p.Uid(), nil
	}

	// No current id exists. Create one.
	uid, err := allocateUniqueUid(xid, instanceIdx, numInstances)
	if err != nil {
		return 0, err
	}

	t := x.DirectedEdge{
		ValueId:   uid,
		Source:    "_assigner_",
		Timestamp: time.Now(),
	}
	_, rerr := pl.AddMutation(context.Background(), t, posting.Set)
	return uid, rerr
}

func StringKey(xid string) []byte {
	return []byte("_uid_|" + xid)
}

// Get returns the uid of the corresponding xid.
func Get(xid string) (uid uint64, rerr error) {
	key := StringKey(xid)
	pl := posting.GetOrCreate(key, uidStore)
	if pl.Length() == 0 {
		return 0, fmt.Errorf("xid: %v doesn't have any uid assigned.", xid)
	}
	if pl.Length() > 1 {
		log.Fatalf("We shouldn't have more than 1 uid for xid: %v\n", xid)
	}
	var p types.Posting
	if ok := pl.Get(&p, 0); !ok {
		return 0, fmt.Errorf("While retrieving entry from posting list")
	}
	return p.Uid(), nil
}

// GetOrAssign returns a unique integer (uid) for a given xid if
// it already exists or assigns a new uid and returns it.
func GetOrAssign(xid string, instanceIdx uint64,
	numInstances uint64) (uid uint64, rerr error) {

	// Prefix _new_ requires us to create a new uid(entity).
	if strings.HasPrefix(xid, "_new_:") {
		return allocateUniqueUid(xid, instanceIdx, numInstances)
	}

	key := StringKey(xid)
	pl := posting.GetOrCreate(key, uidStore)
	if pl.Length() == 0 {
		return assignNew(pl, xid, instanceIdx, numInstances)
	}

	if pl.Length() > 1 {
		log.Fatalf("We shouldn't have more than 1 uid for xid: %v\n", xid)
	} else {
		// We found one posting.
		var p types.Posting
		if ok := pl.Get(&p, 0); !ok {
			return 0, errors.New("While retrieving entry from posting list")
		}
		return p.Uid(), nil
	}
	return 0, errors.New("Some unhandled route lead me here." +
		" Wake the stupid developer up.")
}

// ExternalId returns the xid of a given uid by reading from the uidstore.
// It returns an error if there is no corresponding xid.
func ExternalId(uid uint64) (xid string, rerr error) {
	key := posting.Key(uid, "_xid_") // uid -> "_xid_" -> xid
	pl := posting.GetOrCreate(key, uidStore)
	if pl.Length() == 0 {
		return "", errors.New("NO external id")
	}

	if pl.Length() > 1 {
		log.Fatalf("This shouldn't be happening. Uid: %v", uid)
		return "", errors.New("Multiple external ids for this uid.")
	}

	var p types.Posting
	if ok := pl.Get(&p, 0); !ok {
		return "", errors.New("While retrieving posting")
	}

	if p.Uid() != math.MaxUint64 {
		log.Fatalf("Value uid must be MaxUint64. Uid: %v", p.Uid())
	}
	return string(p.ValueBytes()), rerr
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/bradfitz/gomemcache/memcache"
	"github.com/montanaflynn/stats"
	"time"
	"encoding/binary"
	"strconv"
	"crypto/rand"
	"runtime"
)

func main() {
	initGRs := runtime.NumGoroutine()
        maxGRs := 1
	for {
		for runtime.NumGoroutine() < initGRs + maxGRs {
			go measure(1000)
		}
	}
}

func measure(n int) {
	mc := memcache.New("169.254.10.0:11211")
	timeToSet := make([]float64, n)
	timeToGet := make([]float64, n)
	for i := 0; i < n; i++ {
		timeToSet[i], timeToGet[i] = measureSetAndGetTime(mc)
		if i > 0 && timeToSet[i] < 0 {
			timeToSet[i] = timeToSet[i - 1]
		}
		if i > 0 && timeToGet[i] < 0 {
			timeToGet[i] = timeToGet[i - 1]
		}
	}
	fmt.Printf("Set min: %.2f, max: %.2f, median: %.2f, 95%%: %.2f, 99%%: %.2f \n",
		stats.Min(timeToSet), stats.Max(timeToSet), stats.Median(timeToSet),
		stats.Percentile(timeToSet, 95), stats.Percentile(timeToSet, 99))
	fmt.Printf("Get min: %.2f, max: %.2f, median: %.2f, 95%%: %.2f, 99%%: %.2f \n",
		stats.Min(timeToGet), stats.Max(timeToGet), stats.Median(timeToGet),
		stats.Percentile(timeToGet, 95), stats.Percentile(timeToGet, 99))
}

func measureSetAndGetTime(mc *memcache.Client) (float64, float64) {

	// create key and value
	k := rand12Chars(1) // 12 bytes
	v := []byte(rand12Chars(85)) // 1020 bytes

	// test setting
	start := time.Now().Nanosecond()
	mc.Set(&memcache.Item{Key: k, Value: v})
	timeToSet := float64((time.Now().Nanosecond() - start)) / 1000.0

	// test getting
	start = time.Now().Nanosecond()
	it, err := mc.Get(k)
	if err != nil {
		fmt.Printf("%v\n", err)
	 	return -1, -1
	}
	timeToGet := float64((time.Now().Nanosecond() - start)) / 1000.0

	// validate the value
	if !bytes.Equal(it.Value, v) {
		fmt.Printf("Wrong value: %v\n", it.Value)
	 	return -1, -1
	}
	return timeToSet, timeToGet
}

func rand12Chars(l int) string {
	var buf bytes.Buffer
	for i := 0; i < l; i++ {
		var n uint64
		binary.Read(rand.Reader, binary.LittleEndian, &n)
		buf.WriteString(strconv.FormatUint(n, 36))
	}
	return buf.String()
}
<commit_msg>Added goroutine support<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/bradfitz/gomemcache/memcache"
	"github.com/montanaflynn/stats"
	"time"
	"encoding/binary"
	"strconv"
	"crypto/rand"
	"runtime"
)

func main() {
	initGRs := runtime.NumGoroutine()
        maxGRs := 1
	for {
		for runtime.NumGoroutine() - initGRs < maxGRs {
			go measure(1000)
		}
	}
}

func measure(n int) {
	mc := memcache.New("169.254.10.0:11211")
	timeToSet := make([]float64, n)
	timeToGet := make([]float64, n)
	for i := 0; i < n; i++ {
		timeToSet[i], timeToGet[i] = measureSetAndGetTime(mc)
		if i > 0 && timeToSet[i] < 0 {
			timeToSet[i] = timeToSet[i - 1]
		}
		if i > 0 && timeToGet[i] < 0 {
			timeToGet[i] = timeToGet[i - 1]
		}
	}
	fmt.Printf("Set min: %.2f, max: %.2f, median: %.2f, 95%%: %.2f, 99%%: %.2f \n",
		stats.Min(timeToSet), stats.Max(timeToSet), stats.Median(timeToSet),
		stats.Percentile(timeToSet, 95), stats.Percentile(timeToSet, 99))
	fmt.Printf("Get min: %.2f, max: %.2f, median: %.2f, 95%%: %.2f, 99%%: %.2f \n",
		stats.Min(timeToGet), stats.Max(timeToGet), stats.Median(timeToGet),
		stats.Percentile(timeToGet, 95), stats.Percentile(timeToGet, 99))
}

func measureSetAndGetTime(mc *memcache.Client) (float64, float64) {

	// create key and value
	k := rand12Chars(1) // 12 bytes
	v := []byte(rand12Chars(85)) // 1020 bytes

	// test setting
	start := time.Now().Nanosecond()
	mc.Set(&memcache.Item{Key: k, Value: v})
	timeToSet := float64((time.Now().Nanosecond() - start)) / 1000.0

	// test getting
	start = time.Now().Nanosecond()
	it, err := mc.Get(k)
	if err != nil {
		fmt.Printf("%v\n", err)
	 	return -1, -1
	}
	timeToGet := float64((time.Now().Nanosecond() - start)) / 1000.0

	// validate the value
	if !bytes.Equal(it.Value, v) {
		fmt.Printf("Wrong value: %v\n", it.Value)
	 	return -1, -1
	}
	return timeToSet, timeToGet
}

func rand12Chars(l int) string {
	var buf bytes.Buffer
	for i := 0; i < l; i++ {
		var n uint64
		binary.Read(rand.Reader, binary.LittleEndian, &n)
		buf.WriteString(strconv.FormatUint(n, 36))
	}
	return buf.String()
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/astaxie/beego"
	"runtime"
)

type MessageStruct struct {
	Message string
}

type JsonController struct {
	beego.Controller
}

func (this *JsonController) Get() {
	m := MessageStruct{"Hello, world"}
	this.Data["json"] = &m
	this.ServeJson()
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	beego.Router("/json", &JsonController{})
	beego.Run()
}
<commit_msg>Update hello.go<commit_after>package main

import (
	"github.com/astaxie/beego"
	"runtime"
)

type MessageStruct struct {
	Message string
}

type JsonController struct {
	beego.Controller
}

func (this *JsonController) Get() {
	m := MessageStruct{"Hello, world"}
	this.Data["json"] = &m
	this.ServeJson()
}

func main() {
	//don't need this set, beego default set it
	//runtime.GOMAXPROCS(runtime.NumCPU())
	beego.RunMode = "prod"
	beego.Router("/json", &JsonController{})
	beego.Run()
}
<|endoftext|>


<commit_before>package metadata

import (
	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	. "github.com/cloudfoundry-incubator/docker-cache-acceptance-tests/commons"

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

var _ = Describe("Diego Docker Metadata", func() {
	var appName string

	curlingFunc := func(appName, path string) func() string {
		return func() string {
			return helpers.CurlApp(appName, path)
		}
	}

	Context("with application that has custom docker build instructions", func() {
		BeforeEach(func() {
			appName = generator.RandomName()

			Eventually(cf.Cf("docker-push", appName, "cloudfoundry/diego-docker-app-custom:latest", "--no-start")).Should(Exit(0))
		})

		AfterEach(func() {
			Eventually(cf.Cf("logs", appName, "--recent")).Should(Exit())
			Eventually(cf.Cf("delete", appName, "-f")).Should(Exit(0))
		})

		It("starts", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(helpers.CurlingAppRoot(appName)).Should(ContainSubstring(OK_RESPONSE))
		})

		It("listens on custom port", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(curlingFunc(appName, "/env")).Should(ContainSubstring(`"PORT":"7070"`))
			Consistently(curlingFunc(appName, "/env")).Should(MatchRegexp(`"CF_INSTANCE_PORTS":"\d+:7070`))
		})

		It("uses a custom user", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(curlingFunc(appName, "/env")).Should(ContainSubstring(`"USER":"dockeruser"`))
		})
	})

	Context("with application that has no USER instruction", func() {
		BeforeEach(func() {
			appName = generator.RandomName()

			Eventually(cf.Cf("docker-push", appName, "cloudfoundry/diego-docker-app:latest", "--no-start")).Should(Exit(0))
		})

		AfterEach(func() {
			Eventually(cf.Cf("logs", appName, "--recent")).Should(Exit())
			Eventually(cf.Cf("delete", appName, "-f")).Should(Exit(0))
		})

		It("uses user root", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(curlingFunc(appName, "/env")).Should(ContainSubstring(`"USER":"root"`))
		})
	})
})
<commit_msg>Metadata test now supports new ports format in env<commit_after>package metadata

import (
	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	. "github.com/cloudfoundry-incubator/docker-cache-acceptance-tests/commons"

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

var _ = Describe("Diego Docker Metadata", func() {
	var appName string

	curlingFunc := func(appName, path string) func() string {
		return func() string {
			return helpers.CurlApp(appName, path)
		}
	}

	Context("with application that has custom docker build instructions", func() {
		BeforeEach(func() {
			appName = generator.RandomName()

			Eventually(cf.Cf("docker-push", appName, "cloudfoundry/diego-docker-app-custom:latest", "--no-start")).Should(Exit(0))
		})

		AfterEach(func() {
			Eventually(cf.Cf("logs", appName, "--recent")).Should(Exit())
			Eventually(cf.Cf("delete", appName, "-f")).Should(Exit(0))
		})

		It("starts", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(helpers.CurlingAppRoot(appName)).Should(ContainSubstring(OK_RESPONSE))
		})

		It("listens on custom port", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(curlingFunc(appName, "/env")).Should(ContainSubstring(`"PORT":"7070"`))
			Consistently(curlingFunc(appName, "/env")).Should(MatchRegexp(`"CF_INSTANCE_PORTS":"\[{\\"external\\":\d+,\\"internal\\":7070`))
		})

		It("uses a custom user", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(curlingFunc(appName, "/env")).Should(ContainSubstring(`"USER":"dockeruser"`))
		})
	})

	Context("with application that has no USER instruction", func() {
		BeforeEach(func() {
			appName = generator.RandomName()

			Eventually(cf.Cf("docker-push", appName, "cloudfoundry/diego-docker-app:latest", "--no-start")).Should(Exit(0))
		})

		AfterEach(func() {
			Eventually(cf.Cf("logs", appName, "--recent")).Should(Exit())
			Eventually(cf.Cf("delete", appName, "-f")).Should(Exit(0))
		})

		It("uses user root", func() {
			Eventually(cf.Cf("start", appName)).Should(Exit(0))
			Consistently(curlingFunc(appName, "/env")).Should(ContainSubstring(`"USER":"root"`))
		})
	})
})
<|endoftext|>


<commit_before>package common

import (
	"github.com/mitchellh/goamz/ec2"
)

// BlockDevice
type BlockDevice struct {
	DeviceName          string `mapstructure:"device_name"`
	VirtualName         string `mapstructure:"virtual_name"`
	SnapshotId          string `mapstructure:"snapshot_id"`
	VolumeType          string `mapstructure:"volume_type"`
	VolumeSize          int64  `mapstructure:"volume_size"`
	DeleteOnTermination bool   `mapstructure:"delete_on_termination"`
	IOPS                int64  `mapstructure:"iops"`
	NoDevice            bool   `mapstructure:"no_device"`
}

type BlockDevices struct {
	AMIMappings    []BlockDevice `mapstructure:"ami_block_device_mappings,squash"`
	LaunchMappings []BlockDevice `mapstructure:"launch_block_device_mappings,squash"`
}

func buildBlockDevices(b []BlockDevice) []ec2.BlockDeviceMapping {
	var blockDevices []ec2.BlockDeviceMapping

	for _, blockDevice := range b {
		blockDevices = append(blockDevices, ec2.BlockDeviceMapping{
			DeviceName:          blockDevice.DeviceName,
			VirtualName:         blockDevice.VirtualName,
			SnapshotId:          blockDevice.SnapshotId,
			VolumeType:          blockDevice.VolumeType,
			VolumeSize:          blockDevice.VolumeSize,
			DeleteOnTermination: blockDevice.DeleteOnTermination,
			IOPS:                blockDevice.IOPS,
			NoDevice:            blockDevice.NoDevice,
		})
	}
	return blockDevices
}

func (b *BlockDevices) BuildAMIDevices() []ec2.BlockDeviceMapping {
	return buildBlockDevices(b.AMIMappings)
}

func (b *BlockDevices) BuildLaunchDevices() []ec2.BlockDeviceMapping {
	return buildBlockDevices(b.LaunchMappings)
}
<commit_msg>builder/amazon/common: don't need squash here<commit_after>package common

import (
	"github.com/mitchellh/goamz/ec2"
)

// BlockDevice
type BlockDevice struct {
	DeviceName          string `mapstructure:"device_name"`
	VirtualName         string `mapstructure:"virtual_name"`
	SnapshotId          string `mapstructure:"snapshot_id"`
	VolumeType          string `mapstructure:"volume_type"`
	VolumeSize          int64  `mapstructure:"volume_size"`
	DeleteOnTermination bool   `mapstructure:"delete_on_termination"`
	IOPS                int64  `mapstructure:"iops"`
	NoDevice            bool   `mapstructure:"no_device"`
}

type BlockDevices struct {
	AMIMappings    []BlockDevice `mapstructure:"ami_block_device_mappings"`
	LaunchMappings []BlockDevice `mapstructure:"launch_block_device_mappings"`
}

func buildBlockDevices(b []BlockDevice) []ec2.BlockDeviceMapping {
	var blockDevices []ec2.BlockDeviceMapping

	for _, blockDevice := range b {
		blockDevices = append(blockDevices, ec2.BlockDeviceMapping{
			DeviceName:          blockDevice.DeviceName,
			VirtualName:         blockDevice.VirtualName,
			SnapshotId:          blockDevice.SnapshotId,
			VolumeType:          blockDevice.VolumeType,
			VolumeSize:          blockDevice.VolumeSize,
			DeleteOnTermination: blockDevice.DeleteOnTermination,
			IOPS:                blockDevice.IOPS,
			NoDevice:            blockDevice.NoDevice,
		})
	}
	return blockDevices
}

func (b *BlockDevices) BuildAMIDevices() []ec2.BlockDeviceMapping {
	return buildBlockDevices(b.AMIMappings)
}

func (b *BlockDevices) BuildLaunchDevices() []ec2.BlockDeviceMapping {
	return buildBlockDevices(b.LaunchMappings)
}
<|endoftext|>


<commit_before>package uritemplates

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

var unreserved = regexp.MustCompile("[^A-Za-z0-9\\-._~]")
var reserved = regexp.MustCompile("[^A-Za-z0-9\\-._~:/?#[\\]@!$&'()*+,;=]")

func pctEncode(original string) (result string) {
	for _, b := range []byte(original) {
		if b < 16 {
			result += fmt.Sprintf("%%0%X", b)
		} else {
			result += fmt.Sprintf("%%%X", b)
		}
	}
	return result
}

func escape(s string, allowReserved bool) string {
	var result string
	if allowReserved {
		result = reserved.ReplaceAllStringFunc(s, pctEncode)
	} else {
		result = unreserved.ReplaceAllStringFunc(s, pctEncode)
	}
	return result
}

type UriTemplate struct {
	raw   string
	parts []templatePart
}

func Parse(rawtemplate string) (template *UriTemplate, err error) {
	template = new(UriTemplate)
	template.raw = rawtemplate
	split := strings.Split(rawtemplate, "{")
	template.parts = make([]templatePart, len(split)*2-1)
	for i, s := range split {
		if i == 0 {
			template.parts[i].raw = s
		} else {
			subsplit := strings.Split(s, "}")
			if len(subsplit) != 2 {
				return nil, errors.New("malformed template")
			}
			expression := subsplit[0]
			template.parts[i*2-1] = parseExpression(expression)
			template.parts[i*2].raw = subsplit[1]
		}
	}
	return template, nil
}

const (
	_          = iota
	SIMPLE     = iota
	PLUS       = iota
	SLASH      = iota
	DOT        = iota
	SEMICOLON  = iota
	QUERY      = iota
	AMPERSAND  = iota
	CROSSHATCH = iota
)

type templatePart struct {
	raw           string
	kind          int
	terms         []templateTerm
	first         string
	sep           string
	named         bool
	ifemp         string
	allowReserved bool
}

type templateTerm struct {
	name     string
	explode  bool
	truncate int
}

func parseExpression(expression string) (result templatePart) {
	switch {
	case strings.HasPrefix(expression, "+"):
		result.kind = PLUS
		result.sep = ","
		result.allowReserved = true
	case strings.HasPrefix(expression, "."):
		result.kind = DOT
		result.first = "."
		result.sep = "."
	case strings.HasPrefix(expression, "/"):
		result.kind = SLASH
		result.first = "/"
		result.sep = "/"
	case strings.HasPrefix(expression, ";"):
		result.kind = SEMICOLON
		result.first = ";"
		result.sep = ";"
		result.named = true
	case strings.HasPrefix(expression, "?"):
		result.kind = QUERY
		result.first = "?"
		result.sep = "&"
		result.named = true
		result.ifemp = "="
	case strings.HasPrefix(expression, "&"):
		result.kind = AMPERSAND
		result.first = "&"
		result.sep = "&"
		result.named = true
		result.ifemp = "="
	case strings.HasPrefix(expression, "#"):
		result.kind = CROSSHATCH
		result.first = "#"
		result.sep = ","
		result.allowReserved = true
	default:
		result.kind = SIMPLE
		result.sep = ","
	}
	if result.kind != SIMPLE {
		expression = expression[1:]
	}
	rawterms := strings.Split(expression, ",")
	result.terms = make([]templateTerm, len(rawterms))
	for i, raw := range rawterms {
		result.terms[i] = parseTerm(raw)
	}
	return result
}

func parseTerm(term string) (result templateTerm) {
	if strings.HasSuffix(term, "*") {
		result.explode = true
		term = term[:len(term)-1]
	}
	split := strings.Split(term, ":")
	if len(split) == 1 {
		result.name = term
	} else if len(split) == 2 {
		result.name = split[0]
		parsed, _ := strconv.ParseInt(split[1], 10, 0)
		result.truncate = int(parsed)
	}
	// else error ?
	return result
}

func (self *UriTemplate) ExpandString(values map[string]interface{}) string {
	raw := ""
	for _, p := range self.parts {
		raw = raw + p.expand(values)
	}
	return raw
}

func (self *templatePart) expand(values map[string]interface{}) string {
	if len(self.raw) > 0 {
		return self.raw
	}
	result := self.first
	for _, term := range self.terms {
		value, exists := values[term.name]
		if !exists {
			continue
		}
		var next string
		switch value.(type) {
		case string:
			v := value.(string)
			next = self.expandString(term, v)
		case []interface{}:
			v := value.([]interface{})
			next = self.expandArray(term, v)
		case map[string]interface{}:
			v := value.(map[string]interface{})
			next = self.expandMap(term, v)
		default:
			v := fmt.Sprintf("%v", value)
			next = self.expandString(term, v)
		}
		if result != self.first {
			result += self.sep
		}
		result += next
	}
	if result == self.first {
		result = ""
	}
	return result
}

func (self *templatePart) expandName(name string, empty bool) (result string) {
	if self.named {
		result = name
		if empty {
			result += self.ifemp
		} else {
			result += "="
		}
	}
	return result
}

func (self *templatePart) expandString(t templateTerm, s string) (result string) {
	if len(s) > t.truncate && t.truncate > 0 {
		s = s[:t.truncate]
	}
	return self.expandName(t.name, len(s) == 0) +
		escape(s, self.allowReserved)
}

func (self *templatePart) expandArray(t templateTerm, a []interface{}) (result string) {
	if !t.explode {
		result = self.expandName(t.name, len(a) == 0)
	}
	for i, v := range a {
		if t.explode && i > 0 {
			result += self.sep
		} else if i > 0 {
			result += ","
		}
		var s string
		switch v.(type) {
		case string:
			s = v.(string)
		default:
			s = fmt.Sprintf("%v", v)
		}
		if len(s) > t.truncate && t.truncate > 0 {
			s = s[:t.truncate]
		}
		if self.named && t.explode {
			result += self.expandName(t.name, len(s) == 0)
		}
		result += escape(s, self.allowReserved)
	}
	return result
}

func (self *templatePart) expandMap(t templateTerm, m map[string]interface{}) (result string) {
	for k, v := range m {
		if t.explode && len(result) > 0 {
			result += self.sep
		} else if len(result) > 0 {
			result += ","
		}
		var s string
		switch v.(type) {
		case string:
			s = v.(string)
		default:
			s = fmt.Sprintf("%v", v)
		}
		if len(s) > t.truncate && t.truncate > 0 {
			s = s[:t.truncate]
		}
		if t.explode {
			result += escape(k, self.allowReserved) +
				"=" + escape(s, self.allowReserved)
		} else {
			result += escape(k, self.allowReserved) +
				"," + escape(s, self.allowReserved)
		}
	}
	if !t.explode {
		result = self.expandName(t.name, len(m) == 0) + result
	}
	return result
}
<commit_msg>Pass all but three negative tests.<commit_after>package uritemplates

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

var unreserved = regexp.MustCompile("[^A-Za-z0-9\\-._~]")
var reserved = regexp.MustCompile("[^A-Za-z0-9\\-._~:/?#[\\]@!$&'()*+,;=]")
var validname = regexp.MustCompile("^([A-Za-z0-9_\\.]|%[0-9A-Fa-f][0-9A-Fa-f])+$")

func pctEncode(original string) (result string) {
	for _, b := range []byte(original) {
		if b < 16 {
			result += fmt.Sprintf("%%0%X", b)
		} else {
			result += fmt.Sprintf("%%%X", b)
		}
	}
	return result
}

func escape(s string, allowReserved bool) string {
	var result string
	if allowReserved {
		result = reserved.ReplaceAllStringFunc(s, pctEncode)
	} else {
		result = unreserved.ReplaceAllStringFunc(s, pctEncode)
	}
	return result
}

type UriTemplate struct {
	raw   string
	parts []templatePart
}

func Parse(rawtemplate string) (template *UriTemplate, err error) {
	template = new(UriTemplate)
	template.raw = rawtemplate
	split := strings.Split(rawtemplate, "{")
	template.parts = make([]templatePart, len(split)*2-1)
	for i, s := range split {
		if i == 0 {
			if strings.Contains(s, "}") {
				err = errors.New("unexpected }")
				break
			}
			template.parts[i].raw = s
		} else {
			subsplit := strings.Split(s, "}")
			if len(subsplit) != 2 {
				err = errors.New("malformed template")
				break
			}
			expression := subsplit[0]
			template.parts[i*2-1], err = parseExpression(expression)
			if err != nil { break }
			if strings.Contains(subsplit[1], "}") {
				err = errors.New("unexpected }")
				break
			}
			template.parts[i*2].raw = subsplit[1]
		}
	}
	if err != nil { template = nil }
	return template, err
}

const (
	_          = iota
	SIMPLE     = iota
	PLUS       = iota
	SLASH      = iota
	DOT        = iota
	SEMICOLON  = iota
	QUERY      = iota
	AMPERSAND  = iota
	CROSSHATCH = iota
)

type templatePart struct {
	raw           string
	kind          int
	terms         []templateTerm
	first         string
	sep           string
	named         bool
	ifemp         string
	allowReserved bool
}

type templateTerm struct {
	name     string
	explode  bool
	truncate int
}

func parseExpression(expression string) (result templatePart, err error) {
	switch {
	case strings.HasPrefix(expression, "+"):
		result.kind = PLUS
		result.sep = ","
		result.allowReserved = true
	case strings.HasPrefix(expression, "."):
		result.kind = DOT
		result.first = "."
		result.sep = "."
	case strings.HasPrefix(expression, "/"):
		result.kind = SLASH
		result.first = "/"
		result.sep = "/"
	case strings.HasPrefix(expression, ";"):
		result.kind = SEMICOLON
		result.first = ";"
		result.sep = ";"
		result.named = true
	case strings.HasPrefix(expression, "?"):
		result.kind = QUERY
		result.first = "?"
		result.sep = "&"
		result.named = true
		result.ifemp = "="
	case strings.HasPrefix(expression, "&"):
		result.kind = AMPERSAND
		result.first = "&"
		result.sep = "&"
		result.named = true
		result.ifemp = "="
	case strings.HasPrefix(expression, "#"):
		result.kind = CROSSHATCH
		result.first = "#"
		result.sep = ","
		result.allowReserved = true
	default:
		result.kind = SIMPLE
		result.sep = ","
	}
	if result.kind != SIMPLE {
		expression = expression[1:]
	}
	rawterms := strings.Split(expression, ",")
	result.terms = make([]templateTerm, len(rawterms))
	for i, raw := range rawterms {
		result.terms[i], err = parseTerm(raw)
		if err != nil { break }
	}
	return result, err
}

func parseTerm(term string) (result templateTerm, err error) {
	if strings.HasSuffix(term, "*") {
		result.explode = true
		term = term[:len(term)-1]
	}
	split := strings.Split(term, ":")
	if len(split) == 1 {
		result.name = term
	} else if len(split) == 2 {
		result.name = split[0]
		var parsed int64
		parsed, err = strconv.ParseInt(split[1], 10, 0)
		result.truncate = int(parsed)
	} else {
		err = errors.New("multiple colons in same term")
	}
	if !validname.MatchString(result.name) {
		err = errors.New("not a valid name: " + result.name)
	}
	if result.explode && result.truncate > 0 {
		err = errors.New("both explode and prefix modifers on same term")
	}
	return result, err
}

func (self *UriTemplate) ExpandString(values map[string]interface{}) string {
	raw := ""
	for _, p := range self.parts {
		raw = raw + p.expand(values)
	}
	return raw
}

func (self *templatePart) expand(values map[string]interface{}) string {
	if len(self.raw) > 0 {
		return self.raw
	}
	result := self.first
	for _, term := range self.terms {
		value, exists := values[term.name]
		if !exists {
			continue
		}
		var next string
		switch value.(type) {
		case string:
			v := value.(string)
			next = self.expandString(term, v)
		case []interface{}:
			v := value.([]interface{})
			next = self.expandArray(term, v)
		case map[string]interface{}:
			v := value.(map[string]interface{})
			next = self.expandMap(term, v)
		default:
			v := fmt.Sprintf("%v", value)
			next = self.expandString(term, v)
		}
		if result != self.first {
			result += self.sep
		}
		result += next
	}
	if result == self.first {
		result = ""
	}
	return result
}

func (self *templatePart) expandName(name string, empty bool) (result string) {
	if self.named {
		result = name
		if empty {
			result += self.ifemp
		} else {
			result += "="
		}
	}
	return result
}

func (self *templatePart) expandString(t templateTerm, s string) (result string) {
	if len(s) > t.truncate && t.truncate > 0 {
		s = s[:t.truncate]
	}
	return self.expandName(t.name, len(s) == 0) +
		escape(s, self.allowReserved)
}

func (self *templatePart) expandArray(t templateTerm, a []interface{}) (result string) {
	if !t.explode {
		result = self.expandName(t.name, len(a) == 0)
	}
	for i, v := range a {
		if t.explode && i > 0 {
			result += self.sep
		} else if i > 0 {
			result += ","
		}
		var s string
		switch v.(type) {
		case string:
			s = v.(string)
		default:
			s = fmt.Sprintf("%v", v)
		}
		if len(s) > t.truncate && t.truncate > 0 {
			s = s[:t.truncate]
		}
		if self.named && t.explode {
			result += self.expandName(t.name, len(s) == 0)
		}
		result += escape(s, self.allowReserved)
	}
	return result
}

func (self *templatePart) expandMap(t templateTerm, m map[string]interface{}) (result string) {
	for k, v := range m {
		if t.explode && len(result) > 0 {
			result += self.sep
		} else if len(result) > 0 {
			result += ","
		}
		var s string
		switch v.(type) {
		case string:
			s = v.(string)
		default:
			s = fmt.Sprintf("%v", v)
		}
		if len(s) > t.truncate && t.truncate > 0 {
			s = s[:t.truncate]
		}
		if t.explode {
			result += escape(k, self.allowReserved) +
				"=" + escape(s, self.allowReserved)
		} else {
			result += escape(k, self.allowReserved) +
				"," + escape(s, self.allowReserved)
		}
	}
	if !t.explode {
		result = self.expandName(t.name, len(m) == 0) + result
	}
	return result
}
<|endoftext|>


<commit_before>package digitalocean

import (
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/digitalocean/godo"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
)

type stepSnapshot struct{}

func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
	client := state.Get("client").(*godo.Client)
	ui := state.Get("ui").(packer.Ui)
	c := state.Get("config").(Config)
	dropletId := state.Get("droplet_id").(int)

	ui.Say(fmt.Sprintf("Creating snapshot: %v", c.SnapshotName))
	_, _, err := client.DropletActions.Snapshot(dropletId, c.SnapshotName)
	if err != nil {
		err := fmt.Errorf("Error creating snapshot: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Wait for the droplet to become unlocked first. For snapshots
	// this can end up taking quite a long time, so we hardcode this to
	// 15 minutes.
	if err := waitForDropletUnlocked(client, dropletId, 15*time.Minute); err != nil {
		// If we get an error the first time, actually report it
		err := fmt.Errorf("Error shutting down droplet: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// With the pending state over, verify that we're in the active state
	ui.Say("Waiting for snapshot to complete...")
	err = waitForDropletState("active", dropletId, client, c.StateTimeout)
	if err != nil {
		err := fmt.Errorf("Error waiting for snapshot to complete: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	log.Printf("Looking up snapshot ID for snapshot: %s", c.SnapshotName)
	images, _, err := client.Images.ListUser(&godo.ListOptions{PerPage: 200})
	if err != nil {
		err := fmt.Errorf("Error looking up snapshot ID: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	var imageId int
	for _, image := range images {
		if image.Name == c.SnapshotName {
			imageId = image.ID
			break
		}
	}

	if imageId == 0 {
		err := errors.New("Couldn't find snapshot to get the image ID. Bug?")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	log.Printf("Snapshot image ID: %d", imageId)
	state.Put("snapshot_image_id", imageId)
	state.Put("snapshot_name", c.SnapshotName)
	state.Put("region", c.Region)

	return multistep.ActionContinue
}

func (s *stepSnapshot) Cleanup(state multistep.StateBag) {
	// no cleanup
}
<commit_msg>actually make it 20 mins<commit_after>package digitalocean

import (
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/digitalocean/godo"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
)

type stepSnapshot struct{}

func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
	client := state.Get("client").(*godo.Client)
	ui := state.Get("ui").(packer.Ui)
	c := state.Get("config").(Config)
	dropletId := state.Get("droplet_id").(int)

	ui.Say(fmt.Sprintf("Creating snapshot: %v", c.SnapshotName))
	_, _, err := client.DropletActions.Snapshot(dropletId, c.SnapshotName)
	if err != nil {
		err := fmt.Errorf("Error creating snapshot: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Wait for the droplet to become unlocked first. For snapshots
	// this can end up taking quite a long time, so we hardcode this to
	// 20 minutes.
	if err := waitForDropletUnlocked(client, dropletId, 20*time.Minute); err != nil {
		// If we get an error the first time, actually report it
		err := fmt.Errorf("Error shutting down droplet: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// With the pending state over, verify that we're in the active state
	ui.Say("Waiting for snapshot to complete...")
	err = waitForDropletState("active", dropletId, client, c.StateTimeout)
	if err != nil {
		err := fmt.Errorf("Error waiting for snapshot to complete: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	log.Printf("Looking up snapshot ID for snapshot: %s", c.SnapshotName)
	images, _, err := client.Images.ListUser(&godo.ListOptions{PerPage: 200})
	if err != nil {
		err := fmt.Errorf("Error looking up snapshot ID: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	var imageId int
	for _, image := range images {
		if image.Name == c.SnapshotName {
			imageId = image.ID
			break
		}
	}

	if imageId == 0 {
		err := errors.New("Couldn't find snapshot to get the image ID. Bug?")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	log.Printf("Snapshot image ID: %d", imageId)
	state.Put("snapshot_image_id", imageId)
	state.Put("snapshot_name", c.SnapshotName)
	state.Put("region", c.Region)

	return multistep.ActionContinue
}

func (s *stepSnapshot) Cleanup(state multistep.StateBag) {
	// no cleanup
}
<|endoftext|>