<commit_before>package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/resourced/resourced/libstring"
)

// NewDefaultConfigs provide default config setup.
// This function is called on first boot.
func NewDefaultConfigs(configDir string) error {
	configDir = libstring.ExpandTildeAndEnv(configDir)

	// Create configDir if it does not exist
	if _, err := os.Stat(configDir); err != nil {
		if os.IsNotExist(err) {
			err := os.MkdirAll(configDir, 0755)
			if err != nil {
				return err
			}

			logrus.WithFields(logrus.Fields{
				"Directory": configDir,
			}).Infof("Created config directory")
		}
	}

	// Create subdirectories
	for _, subdirConfigs := range []string{"readers", "writers", "executors", "tags"} {
		subdirPath := path.Join(configDir, subdirConfigs)

		if _, err := os.Stat(subdirPath); err != nil {
			if os.IsNotExist(err) {
				err := os.MkdirAll(subdirPath, 0755)
				if err != nil {
					return err
				}

				logrus.WithFields(logrus.Fields{
					"Directory": subdirPath,
				}).Infof("Created config directory")
			}
		}
	}

	// TODO(didip): Download default readers
	// https://raw.githubusercontent.com/resourced/resourced/master/tests/data/resourced-configs/readers/cpu-info.toml

	// Create default tags
	defaultTagsTemplate := `GOOS=%v
uname=%v
`

	unameBytes, err := exec.Command("uname", "-a").CombinedOutput()
	if err != nil {
		return err
	}

	uname := strings.TrimSpace(string(unameBytes))

	err = ioutil.WriteFile(path.Join(configDir, "tags", "default"), []byte(fmt.Sprintf(defaultTagsTemplate, runtime.GOOS, uname)), 0755)
	if err != nil {
		return err
	}
	logrus.WithFields(logrus.Fields{
		"File": path.Join(configDir, "tags", "default"),
	}).Infof("Created default tags file")

	// Create a default general.toml
	generalToml := `# Addr is the host and port of ResourceD Agent HTTP/S server
Addr = "localhost:55555"

# Valid LogLevel are: debug, info, warning, error, fatal, panic
LogLevel = "info"

[HTTPS]
CertFile = ""
KeyFile = ""

[ResourcedMaster]
# Url is the root endpoint to Resourced Master
URL = "http://localhost:55655"

# General purpose AccessToken, it will be used when AccessToken is not defined elsewhere.
AccessToken = "{access-token}"
`

	err = ioutil.WriteFile(path.Join(configDir, "general.toml"), []byte(generalToml), 0644)
	if err != nil {
		return err
	}
	logrus.WithFields(logrus.Fields{
		"File": path.Join(configDir, "general.toml"),
	}).Infof("Created general config file")

	return nil
}
<commit_msg>Pull down default reader config files on first boot.<commit_after>package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/resourced/resourced/libstring"
)

// NewDefaultConfigs provide default config setup.
// This function is called on first boot.
func NewDefaultConfigs(configDir string) error {
	configDir = libstring.ExpandTildeAndEnv(configDir)

	// Create configDir if it does not exist
	if _, err := os.Stat(configDir); err != nil {
		if os.IsNotExist(err) {
			err := os.MkdirAll(configDir, 0755)
			if err != nil {
				return err
			}

			logrus.WithFields(logrus.Fields{
				"Directory": configDir,
			}).Infof("Created config directory")
		}
	}

	// Create subdirectories
	for _, subdirConfigs := range []string{"readers", "writers", "executors", "tags"} {
		subdirPath := path.Join(configDir, subdirConfigs)

		if _, err := os.Stat(subdirPath); err != nil {
			if os.IsNotExist(err) {
				err := os.MkdirAll(subdirPath, 0755)
				if err != nil {
					return err
				}

				logrus.WithFields(logrus.Fields{
					"Directory": subdirPath,
				}).Infof("Created config directory")

				// Download default reader config files
				// Ignore errors as it's not important.
				if subdirConfigs == "readers" {
					output, err := exec.Command(
						"svn", "checkout",
						"https://github.com/resourced/resourced/trunk/tests/data/resourced-configs/readers",
						subdirPath,
					).CombinedOutput()

					if err != nil {
						logrus.WithFields(logrus.Fields{
							"Error": err.Error(),
						}).Error("Failed to download default reader config files: " + string(output))
					}

					// Remove .svn artifacts
					os.RemoveAll(path.Join(subdirPath, ".svn"))
				}
			}
		}
	}

	// Create default tags
	defaultTagsTemplate := `GOOS=%v
uname=%v
`
	unameBytes, err := exec.Command("uname", "-a").CombinedOutput()
	if err != nil {
		return err
	}

	uname := strings.TrimSpace(string(unameBytes))

	err = ioutil.WriteFile(path.Join(configDir, "tags", "default"), []byte(fmt.Sprintf(defaultTagsTemplate, runtime.GOOS, uname)), 0755)
	if err != nil {
		return err
	}
	logrus.WithFields(logrus.Fields{
		"File": path.Join(configDir, "tags", "default"),
	}).Infof("Created default tags file")

	// Create a default general.toml
	generalToml := `# Addr is the host and port of ResourceD Agent HTTP/S server
Addr = "localhost:55555"

# Valid LogLevel are: debug, info, warning, error, fatal, panic
LogLevel = "info"

[HTTPS]
CertFile = ""
KeyFile = ""

[ResourcedMaster]
# Url is the root endpoint to Resourced Master
URL = "http://localhost:55655"

# General purpose AccessToken, it will be used when AccessToken is not defined elsewhere.
AccessToken = "{access-token}"
`

	err = ioutil.WriteFile(path.Join(configDir, "general.toml"), []byte(generalToml), 0644)
	if err != nil {
		return err
	}
	logrus.WithFields(logrus.Fields{
		"File": path.Join(configDir, "general.toml"),
	}).Infof("Created general config file")

	return nil
}
<|endoftext|>


<commit_before>package memory

import (
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"strings"
)

// the collision avoidance window defines how many times we try to find a higher
// slot that's free if two record hashes collide
var collisionAvoidanceWindow = uint32(1024)

// the function we use to get the hash for hashing the meta records
// it can be replaced for mocking in tests
var queryHash func() hash.Hash32

func init() {
	queryHash = fnv.New32a
}

type metaTagRecord struct {
	metaTags []kv
	queries  []expression
}

type recordId uint32

// list of meta records keyed by a unique identifier used as ID
type metaTagRecords map[recordId]metaTagRecord

// index structure keyed by tag -> value -> list of meta record IDs
type metaTagValue map[string][]recordId
type metaTagIndex map[string]metaTagValue

func (m metaTagIndex) deleteRecord(keyValue kv, id recordId) {
	if values, ok := m[keyValue.key]; ok {
		if ids, ok := values[keyValue.value]; ok {
			for i := 0; i < len(ids); i++ {
				if ids[i] == id {
					// no need to keep the order
					ids[i] = ids[len(ids)-1]
					values[keyValue.value] = ids[:len(ids)-1]

					// no id should ever be present more than once
					return
				}
			}
		}
	}
}

func (m metaTagIndex) insertRecord(keyValue kv, id recordId) {
	var values metaTagValue
	var ok bool

	if values, ok = m[keyValue.key]; !ok {
		values = make(metaTagValue)
		m[keyValue.key] = values
	}

	values[keyValue.value] = append(values[keyValue.value], id)
}

// metaTagRecordFromStrings takes two slices of strings, parses them and returns a metaTagRecord
// The first slice of strings has the meta tags & values
// The second slice has the tag query expressions which the meta tags & values refer to
// On parsing error the second returned value is an error, otherwise it is nil
func metaTagRecordFromStrings(metaTags []string, tagQueryExpressions []string) (metaTagRecord, error) {
	record := metaTagRecord{
		metaTags: make([]kv, 0, len(metaTags)),
		queries:  make([]expression, 0, len(tagQueryExpressions)),
	}

	if len(tagQueryExpressions) == 0 {
		return record, fmt.Errorf("Requiring at least one tag query expression, 0 given")
	}

	for _, tag := range metaTags {
		tagSplits := strings.SplitN(tag, "=", 2)
		if len(tagSplits) < 2 {
			return record, fmt.Errorf("Missing \"=\" sign in tag %s", tag)
		}
		if len(tagSplits[0]) == 0 || len(tagSplits[1]) == 0 {
			return record, fmt.Errorf("Tag/Value cannot be empty in %s", tag)
		}

		record.metaTags = append(record.metaTags, kv{key: tagSplits[0], value: tagSplits[1]})
	}

	for _, query := range tagQueryExpressions {
		parsed, err := parseExpression(query)
		if err != nil {
			return record, err
		}
		record.queries = append(record.queries, parsed)
	}

	return record, nil
}

func (m *metaTagRecord) metaTagStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.metaTags))

	for i, tag := range m.metaTags {
		tag.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

func (m *metaTagRecord) queryStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.queries))

	for i, query := range m.queries {
		query.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

// hashQueries generates a hash of all the queries in the record
func (m *metaTagRecord) hashQueries() recordId {
	builder := strings.Builder{}
	for _, query := range m.queries {
		query.stringIntoBuilder(&builder)

		// trailing ";" doesn't matter, this is only hash input
		builder.WriteString(";")
	}

	h := queryHash()
	h.Write([]byte(builder.String()))
	return recordId(h.Sum32())
}

// sortQueries sorts all the queries first by key, then by value, then by
// operator. The order doesn't matter, it only needs to be consistent
func (m *metaTagRecord) sortQueries() {
	sort.Slice(m.queries, func(i, j int) bool {
		if m.queries[i].key == m.queries[j].key {
			if m.queries[i].value == m.queries[j].value {
				return m.queries[i].operator < m.queries[j].operator
			}
			return m.queries[i].value < m.queries[j].value
		}
		return m.queries[i].key < m.queries[j].key
	})
}

// matchesQueries compares another tag record's queries to this
// one's queries. Returns true if they are equal, otherwise false.
// It is assumed that all the queries are already sorted
func (m *metaTagRecord) matchesQueries(other metaTagRecord) bool {
	if len(m.queries) != len(other.queries) {
		return false
	}

	for i, query := range m.queries {
		if query.key != other.queries[i].key {
			return false
		}

		if query.operator != other.queries[i].operator {
			return false
		}

		if query.value != other.queries[i].value {
			return false
		}
	}

	return true
}

// hasMetaTags returns true if the meta tag record has one or more
// meta tags, otherwise it returns false
func (m *metaTagRecord) hasMetaTags() bool {
	return len(m.metaTags) > 0
}

// upsert inserts or updates a meta tag record according to the given specifications
// it uses the set of tag query expressions as the identity of the record, if a record with the
// same identity is already present then its meta tags get updated to the specified ones.
// If the new record contains no meta tags, then the update is equivalent to a delete.
// Those are the return values:
// 1) The id at which the new record got inserted
// 2) Pointer to the inserted metaTagRecord
// 3) The id of the record that has been replaced if an update was performed
// 4) Pointer to the metaTagRecord that has been replaced if an update was performed, otherwise nil
// 5) Error if an error occurred, otherwise it's nil
func (m metaTagRecords) upsert(metaTags []string, tagQueryExpressions []string) (recordId, *metaTagRecord, recordId, *metaTagRecord, error) {
	record, err := metaTagRecordFromStrings(metaTags, tagQueryExpressions)
	if err != nil {
		return 0, nil, 0, nil, err
	}

	record.sortQueries()
	id := record.hashQueries()
	var oldRecord *metaTagRecord
	var oldId recordId

	// loop over existing records, starting from id, trying to find one that has
	// the exact same queries as the one we're upserting
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		if existingRecord, ok := m[id+recordId(i)]; ok {
			if record.matchesQueries(existingRecord) {
				oldRecord = &existingRecord
				oldId = id + recordId(i)
				delete(m, oldId)
				break
			}
		}
	}

	if !record.hasMetaTags() {
		return 0, &record, oldId, oldRecord, nil
	}

	// now find the best position to insert the new/updated record, starting from id
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		// if we find a free slot, then insert the new record there
		if _, ok := m[id]; !ok {
			m[id] = record

			return id, &record, oldId, oldRecord, nil
		}

		id++
	}

	return 0, nil, 0, nil, fmt.Errorf("MetaTagRecordUpsert: Unable to find a slot to insert record")
}
<commit_msg>better comment<commit_after>package memory

import (
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"strings"
)

// the collision avoidance window defines how many times we try to find a higher
// slot that's free if two record hashes collide
var collisionAvoidanceWindow = uint32(1024)

// the function we use to get the hash for hashing the meta records
// it can be replaced for mocking in tests
var queryHash func() hash.Hash32

func init() {
	queryHash = fnv.New32a
}

type metaTagRecord struct {
	metaTags []kv
	queries  []expression
}

type recordId uint32

// list of meta records keyed by a unique identifier used as ID
type metaTagRecords map[recordId]metaTagRecord

// index structure keyed by tag -> value -> list of meta record IDs
type metaTagValue map[string][]recordId
type metaTagIndex map[string]metaTagValue

func (m metaTagIndex) deleteRecord(keyValue kv, id recordId) {
	if values, ok := m[keyValue.key]; ok {
		if ids, ok := values[keyValue.value]; ok {
			for i := 0; i < len(ids); i++ {
				if ids[i] == id {
					// no need to keep the order
					ids[i] = ids[len(ids)-1]
					values[keyValue.value] = ids[:len(ids)-1]

					// no id should ever be present more than once
					return
				}
			}
		}
	}
}

func (m metaTagIndex) insertRecord(keyValue kv, id recordId) {
	var values metaTagValue
	var ok bool

	if values, ok = m[keyValue.key]; !ok {
		values = make(metaTagValue)
		m[keyValue.key] = values
	}

	values[keyValue.value] = append(values[keyValue.value], id)
}

// metaTagRecordFromStrings takes two slices of strings, parses them and returns a metaTagRecord
// The first slice of strings has the meta tags & values
// The second slice has the tag query expressions which the meta tags & values refer to
// On parsing error the second returned value is an error, otherwise it is nil
func metaTagRecordFromStrings(metaTags []string, tagQueryExpressions []string) (metaTagRecord, error) {
	record := metaTagRecord{
		metaTags: make([]kv, 0, len(metaTags)),
		queries:  make([]expression, 0, len(tagQueryExpressions)),
	}

	if len(tagQueryExpressions) == 0 {
		return record, fmt.Errorf("Requiring at least one tag query expression, 0 given")
	}

	for _, tag := range metaTags {
		tagSplits := strings.SplitN(tag, "=", 2)
		if len(tagSplits) < 2 {
			return record, fmt.Errorf("Missing \"=\" sign in tag %s", tag)
		}
		if len(tagSplits[0]) == 0 || len(tagSplits[1]) == 0 {
			return record, fmt.Errorf("Tag/Value cannot be empty in %s", tag)
		}

		record.metaTags = append(record.metaTags, kv{key: tagSplits[0], value: tagSplits[1]})
	}

	for _, query := range tagQueryExpressions {
		parsed, err := parseExpression(query)
		if err != nil {
			return record, err
		}
		record.queries = append(record.queries, parsed)
	}

	return record, nil
}

func (m *metaTagRecord) metaTagStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.metaTags))

	for i, tag := range m.metaTags {
		tag.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

func (m *metaTagRecord) queryStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.queries))

	for i, query := range m.queries {
		query.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

// hashQueries generates a hash of all the queries in the record
func (m *metaTagRecord) hashQueries() recordId {
	builder := strings.Builder{}
	for _, query := range m.queries {
		query.stringIntoBuilder(&builder)

		// trailing ";" doesn't matter, this is only hash input
		builder.WriteString(";")
	}

	h := queryHash()
	h.Write([]byte(builder.String()))
	return recordId(h.Sum32())
}

// sortQueries sorts all the queries first by key, then by value, then by
// operator. The order itself is not relevant, we only need to guarantee
// that the order is consistent for a given set of queries after every call
// to this function, because the queries will then be used as hash input
func (m *metaTagRecord) sortQueries() {
	sort.Slice(m.queries, func(i, j int) bool {
		if m.queries[i].key == m.queries[j].key {
			if m.queries[i].value == m.queries[j].value {
				return m.queries[i].operator < m.queries[j].operator
			}
			return m.queries[i].value < m.queries[j].value
		}
		return m.queries[i].key < m.queries[j].key
	})
}

// matchesQueries compares another tag record's queries to this
// one's queries. Returns true if they are equal, otherwise false.
// It is assumed that all the queries are already sorted
func (m *metaTagRecord) matchesQueries(other metaTagRecord) bool {
	if len(m.queries) != len(other.queries) {
		return false
	}

	for i, query := range m.queries {
		if query.key != other.queries[i].key {
			return false
		}

		if query.operator != other.queries[i].operator {
			return false
		}

		if query.value != other.queries[i].value {
			return false
		}
	}

	return true
}

// hasMetaTags returns true if the meta tag record has one or more
// meta tags, otherwise it returns false
func (m *metaTagRecord) hasMetaTags() bool {
	return len(m.metaTags) > 0
}

// upsert inserts or updates a meta tag record according to the given specifications
// it uses the set of tag query expressions as the identity of the record, if a record with the
// same identity is already present then its meta tags get updated to the specified ones.
// If the new record contains no meta tags, then the update is equivalent to a delete.
// Those are the return values:
// 1) The id at which the new record got inserted
// 2) Pointer to the inserted metaTagRecord
// 3) The id of the record that has been replaced if an update was performed
// 4) Pointer to the metaTagRecord that has been replaced if an update was performed, otherwise nil
// 5) Error if an error occurred, otherwise it's nil
func (m metaTagRecords) upsert(metaTags []string, tagQueryExpressions []string) (recordId, *metaTagRecord, recordId, *metaTagRecord, error) {
	record, err := metaTagRecordFromStrings(metaTags, tagQueryExpressions)
	if err != nil {
		return 0, nil, 0, nil, err
	}

	record.sortQueries()
	id := record.hashQueries()
	var oldRecord *metaTagRecord
	var oldId recordId

	// loop over existing records, starting from id, trying to find one that has
	// the exact same queries as the one we're upserting
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		if existingRecord, ok := m[id+recordId(i)]; ok {
			if record.matchesQueries(existingRecord) {
				oldRecord = &existingRecord
				oldId = id + recordId(i)
				delete(m, oldId)
				break
			}
		}
	}

	if !record.hasMetaTags() {
		return 0, &record, oldId, oldRecord, nil
	}

	// now find the best position to insert the new/updated record, starting from id
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		// if we find a free slot, then insert the new record there
		if _, ok := m[id]; !ok {
			m[id] = record

			return id, &record, oldId, oldRecord, nil
		}

		id++
	}

	return 0, nil, 0, nil, fmt.Errorf("MetaTagRecordUpsert: Unable to find a slot to insert record")
}
<|endoftext|>


<commit_before>package decor

import (
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/VividCortex/ewma"
)

type SpeedKiB float64

func (s SpeedKiB) Format(st fmt.State, verb rune) {
	prec, ok := st.Precision()

	if verb == 'd' || !ok {
		prec = 0
	}
	if verb == 'f' && !ok {
		prec = 6
	}
	// retain old beahavior if s verb used
	if verb == 's' {
		prec = 1
	}

	var res, unit string
	switch {
	case s >= TiB:
		unit = "TiB/s"
		res = strconv.FormatFloat(float64(s)/TiB, 'f', prec, 64)
	case s >= GiB:
		unit = "GiB/s"
		res = strconv.FormatFloat(float64(s)/GiB, 'f', prec, 64)
	case s >= MiB:
		unit = "MiB/s"
		res = strconv.FormatFloat(float64(s)/MiB, 'f', prec, 64)
	case s >= KiB:
		unit = "KiB/s"
		res = strconv.FormatFloat(float64(s)/KiB, 'f', prec, 64)
	default:
		unit = "b/s"
		res = strconv.FormatInt(int64(s), 10)
	}

	if st.Flag(' ') {
		res += " "
	}
	res += unit

	if w, ok := st.Width(); ok {
		if len(res) < w {
			pad := strings.Repeat(" ", w-len(res))
			if st.Flag(int('-')) {
				res += pad
			} else {
				res = pad + res
			}
		}
	}

	io.WriteString(st, res)
}

type SpeedKB float64

func (s SpeedKB) Format(st fmt.State, verb rune) {
	prec, ok := st.Precision()

	if verb == 'd' || !ok {
		prec = 0
	}
	if verb == 'f' && !ok {
		prec = 6
	}
	// retain old beahavior if s verb used
	if verb == 's' {
		prec = 1
	}

	var res, unit string
	switch {
	case s >= TB:
		unit = "TB/s"
		res = strconv.FormatFloat(float64(s)/TB, 'f', prec, 64)
	case s >= GB:
		unit = "GB/s"
		res = strconv.FormatFloat(float64(s)/GB, 'f', prec, 64)
	case s >= MB:
		unit = "MB/s"
		res = strconv.FormatFloat(float64(s)/MB, 'f', prec, 64)
	case s >= KB:
		unit = "kB/s"
		res = strconv.FormatFloat(float64(s)/KB, 'f', prec, 64)
	default:
		unit = "b/s"
		res = strconv.FormatInt(int64(s), 10)
	}

	if st.Flag(' ') {
		res += " "
	}
	res += unit

	if w, ok := st.Width(); ok {
		if len(res) < w {
			pad := strings.Repeat(" ", w-len(res))
			if st.Flag(int('-')) {
				res += pad
			} else {
				res = pad + res
			}
		}
	}

	io.WriteString(st, res)
}

// SpeedNoUnit returns raw I/O operation speed decorator.
//
//	`unitFormat` printf compatible verb for value, like "%f" or "%d"
//
//	`age` is the previous N samples to average over.
//	 If zero value provided, it defaults to 30.
//
//	`sbCh` is a start block receive channel. User suppose to send time.Now()
//	 to this channel on each iteration of a start block, right before actual job.
//	 The channel will be auto closed on bar shutdown event, so there is no need
//	 to close from user side.
//
//	`wcc` optional WC config
//
// unitFormat example:
//
//	"%.1f" = "1.0" or "% .1f" = "1.0"
func SpeedNoUnit(unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	return speed(0, unitFormat, age, sbCh, wcc...)
}

// SpeedKibiByte returns human friendly I/O operation speed decorator,
//
//	`unitFormat` printf compatible verb for value, like "%f" or "%d"
//
//	`age` is the previous N samples to average over.
//	 If zero value provided, it defaults to 30.
//
//	`sbCh` is a start block receive channel. User suppose to send time.Now()
//	 to this channel on each iteration of a start block, right before actual job.
//	 The channel will be auto closed on bar shutdown event, so there is no need
//	 to close from user side.
//
//	`wcc` optional WC config
//
// unitFormat example:
//
//	"%.1f" = "1.0MiB/s" or "% .1f" = "1.0 MiB/s"
func SpeedKibiByte(unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	return speed(unitKiB, unitFormat, age, sbCh, wcc...)
}

// SpeedKiloByte returns human friendly I/O operation speed decorator,
//
//	`unitFormat` printf compatible verb for value, like "%f" or "%d"
//
//	`age` is the previous N samples to average over.
//	 If zero value provided, it defaults to 30.
//
//	`sbCh` is a start block receive channel. User suppose to send time.Now()
//	 to this channel on each iteration of a start block, right before actual job.
//	 The channel will be auto closed on bar shutdown event, so there is no need
//	 to close from user side.
//
//	`wcc` optional WC config
//
// unitFormat example:
//
//	"%.1f" = "1.0MB/s" or "% .1f" = "1.0 MB/s"
func SpeedKiloByte(unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	return speed(unitKB, unitFormat, age, sbCh, wcc...)
}

func speed(unit int, unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	if sbCh == nil {
		panic("start block channel must not be nil")
	}
	var wc WC
	for _, widthConf := range wcc {
		wc = widthConf
	}
	wc.BuildFormat()
	if age == .0 {
		age = ewma.AVG_METRIC_AGE
	}
	d := &ewmaSpeed{
		unit:       unit,
		unitFormat: unitFormat,
		wc:         wc,
		mAverage:   ewma.NewMovingAverage(age),
		sbReceiver: sbCh,
		sbStreamer: make(chan time.Time),
	}
	go d.serve()
	return d
}

type ewmaSpeed struct {
	unit       int
	unitFormat string
	wc         WC
	mAverage   ewma.MovingAverage
	sbReceiver chan time.Time
	sbStreamer chan time.Time
	onComplete *struct {
		msg string
		wc  WC
	}
}

func (s *ewmaSpeed) Decor(st *Statistics, widthAccumulator chan<- int, widthDistributor <-chan int) string {
	if st.Completed && s.onComplete != nil {
		return s.onComplete.wc.FormatMsg(s.onComplete.msg, widthAccumulator, widthDistributor)
	}
	var str string
	speed := round(s.mAverage.Value())
	switch s.unit {
	case unitKiB:
		str = fmt.Sprintf(s.unitFormat, SpeedKiB(speed))
	case unitKB:
		str = fmt.Sprintf(s.unitFormat, SpeedKB(speed))
	default:
		str = fmt.Sprintf(s.unitFormat, speed)
	}
	return s.wc.FormatMsg(str, widthAccumulator, widthDistributor)
}

func (s *ewmaSpeed) NextAmount(n int) {
	sb := <-s.sbStreamer
	speed := float64(n) / time.Since(sb).Seconds()
	s.mAverage.Add(speed)
}

func (s *ewmaSpeed) OnCompleteMessage(msg string, wcc ...WC) {
	var wc WC
	for _, widthConf := range wcc {
		wc = widthConf
	}
	wc.BuildFormat()
	s.onComplete = &struct {
		msg string
		wc  WC
	}{msg, wc}
}

func (s *ewmaSpeed) Shutdown() {
	close(s.sbReceiver)
}

func (s *ewmaSpeed) serve() {
	for now := range s.sbReceiver {
		s.sbStreamer <- now
	}
}
<commit_msg>no need to round speed<commit_after>package decor

import (
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/VividCortex/ewma"
)

type SpeedKiB float64

func (s SpeedKiB) Format(st fmt.State, verb rune) {
	prec, ok := st.Precision()

	if verb == 'd' || !ok {
		prec = 0
	}
	if verb == 'f' && !ok {
		prec = 6
	}
	// retain old beahavior if s verb used
	if verb == 's' {
		prec = 1
	}

	var res, unit string
	switch {
	case s >= TiB:
		unit = "TiB/s"
		res = strconv.FormatFloat(float64(s)/TiB, 'f', prec, 64)
	case s >= GiB:
		unit = "GiB/s"
		res = strconv.FormatFloat(float64(s)/GiB, 'f', prec, 64)
	case s >= MiB:
		unit = "MiB/s"
		res = strconv.FormatFloat(float64(s)/MiB, 'f', prec, 64)
	case s >= KiB:
		unit = "KiB/s"
		res = strconv.FormatFloat(float64(s)/KiB, 'f', prec, 64)
	default:
		unit = "b/s"
		res = strconv.FormatInt(int64(s), 10)
	}

	if st.Flag(' ') {
		res += " "
	}
	res += unit

	if w, ok := st.Width(); ok {
		if len(res) < w {
			pad := strings.Repeat(" ", w-len(res))
			if st.Flag(int('-')) {
				res += pad
			} else {
				res = pad + res
			}
		}
	}

	io.WriteString(st, res)
}

type SpeedKB float64

func (s SpeedKB) Format(st fmt.State, verb rune) {
	prec, ok := st.Precision()

	if verb == 'd' || !ok {
		prec = 0
	}
	if verb == 'f' && !ok {
		prec = 6
	}
	// retain old beahavior if s verb used
	if verb == 's' {
		prec = 1
	}

	var res, unit string
	switch {
	case s >= TB:
		unit = "TB/s"
		res = strconv.FormatFloat(float64(s)/TB, 'f', prec, 64)
	case s >= GB:
		unit = "GB/s"
		res = strconv.FormatFloat(float64(s)/GB, 'f', prec, 64)
	case s >= MB:
		unit = "MB/s"
		res = strconv.FormatFloat(float64(s)/MB, 'f', prec, 64)
	case s >= KB:
		unit = "kB/s"
		res = strconv.FormatFloat(float64(s)/KB, 'f', prec, 64)
	default:
		unit = "b/s"
		res = strconv.FormatInt(int64(s), 10)
	}

	if st.Flag(' ') {
		res += " "
	}
	res += unit

	if w, ok := st.Width(); ok {
		if len(res) < w {
			pad := strings.Repeat(" ", w-len(res))
			if st.Flag(int('-')) {
				res += pad
			} else {
				res = pad + res
			}
		}
	}

	io.WriteString(st, res)
}

// SpeedNoUnit returns raw I/O operation speed decorator.
//
//	`unitFormat` printf compatible verb for value, like "%f" or "%d"
//
//	`age` is the previous N samples to average over.
//	 If zero value provided, it defaults to 30.
//
//	`sbCh` is a start block receive channel. User suppose to send time.Now()
//	 to this channel on each iteration of a start block, right before actual job.
//	 The channel will be auto closed on bar shutdown event, so there is no need
//	 to close from user side.
//
//	`wcc` optional WC config
//
// unitFormat example:
//
//	"%.1f" = "1.0" or "% .1f" = "1.0"
func SpeedNoUnit(unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	return speed(0, unitFormat, age, sbCh, wcc...)
}

// SpeedKibiByte returns human friendly I/O operation speed decorator,
//
//	`unitFormat` printf compatible verb for value, like "%f" or "%d"
//
//	`age` is the previous N samples to average over.
//	 If zero value provided, it defaults to 30.
//
//	`sbCh` is a start block receive channel. User suppose to send time.Now()
//	 to this channel on each iteration of a start block, right before actual job.
//	 The channel will be auto closed on bar shutdown event, so there is no need
//	 to close from user side.
//
//	`wcc` optional WC config
//
// unitFormat example:
//
//	"%.1f" = "1.0MiB/s" or "% .1f" = "1.0 MiB/s"
func SpeedKibiByte(unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	return speed(unitKiB, unitFormat, age, sbCh, wcc...)
}

// SpeedKiloByte returns human friendly I/O operation speed decorator,
//
//	`unitFormat` printf compatible verb for value, like "%f" or "%d"
//
//	`age` is the previous N samples to average over.
//	 If zero value provided, it defaults to 30.
//
//	`sbCh` is a start block receive channel. User suppose to send time.Now()
//	 to this channel on each iteration of a start block, right before actual job.
//	 The channel will be auto closed on bar shutdown event, so there is no need
//	 to close from user side.
//
//	`wcc` optional WC config
//
// unitFormat example:
//
//	"%.1f" = "1.0MB/s" or "% .1f" = "1.0 MB/s"
func SpeedKiloByte(unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	return speed(unitKB, unitFormat, age, sbCh, wcc...)
}

func speed(unit int, unitFormat string, age float64, sbCh chan time.Time, wcc ...WC) Decorator {
	if sbCh == nil {
		panic("start block channel must not be nil")
	}
	var wc WC
	for _, widthConf := range wcc {
		wc = widthConf
	}
	wc.BuildFormat()
	if age == .0 {
		age = ewma.AVG_METRIC_AGE
	}
	d := &ewmaSpeed{
		unit:       unit,
		unitFormat: unitFormat,
		wc:         wc,
		mAverage:   ewma.NewMovingAverage(age),
		sbReceiver: sbCh,
		sbStreamer: make(chan time.Time),
	}
	go d.serve()
	return d
}

type ewmaSpeed struct {
	unit       int
	unitFormat string
	wc         WC
	mAverage   ewma.MovingAverage
	sbReceiver chan time.Time
	sbStreamer chan time.Time
	onComplete *struct {
		msg string
		wc  WC
	}
}

func (s *ewmaSpeed) Decor(st *Statistics, widthAccumulator chan<- int, widthDistributor <-chan int) string {
	if st.Completed && s.onComplete != nil {
		return s.onComplete.wc.FormatMsg(s.onComplete.msg, widthAccumulator, widthDistributor)
	}
	var str string
	speed := s.mAverage.Value()
	switch s.unit {
	case unitKiB:
		str = fmt.Sprintf(s.unitFormat, SpeedKiB(speed))
	case unitKB:
		str = fmt.Sprintf(s.unitFormat, SpeedKB(speed))
	default:
		str = fmt.Sprintf(s.unitFormat, speed)
	}
	return s.wc.FormatMsg(str, widthAccumulator, widthDistributor)
}

func (s *ewmaSpeed) NextAmount(n int) {
	sb := <-s.sbStreamer
	speed := float64(n) / time.Since(sb).Seconds()
	s.mAverage.Add(speed)
}

func (s *ewmaSpeed) OnCompleteMessage(msg string, wcc ...WC) {
	var wc WC
	for _, widthConf := range wcc {
		wc = widthConf
	}
	wc.BuildFormat()
	s.onComplete = &struct {
		msg string
		wc  WC
	}{msg, wc}
}

func (s *ewmaSpeed) Shutdown() {
	close(s.sbReceiver)
}

func (s *ewmaSpeed) serve() {
	for now := range s.sbReceiver {
		s.sbStreamer <- now
	}
}
<|endoftext|>


<commit_before>package imagick

import (
	"fmt"
	"io/ioutil"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/pressly/imgry"
)

func TestLoadBlob(t *testing.T) {
	tdImage1, err := ioutil.ReadFile("../testdata/image1.jpg")
	assert.NoError(t, err)

	ng := Engine{}
	im, err := ng.LoadBlob(tdImage1)
	assert.NoError(t, err)
	defer im.Release()

	sz, _ := imgry.NewSizingFromQuery("size=800x")
	err = im.SizeIt(sz)
	assert.NoError(t, err)

	im2Path := "/tmp/image1.jpg"
	im.WriteToFile(im2Path)

	im2, err := ng.LoadFile(im2Path)
	assert.NoError(t, err)

	assert.True(t, im2.Width() == 800)
	assert.Equal(t, "jpg", im2.Format())

	err = im2.SetFormat("png")
	assert.NoError(t, err)
}

func TestGetImageInfo(t *testing.T) {
	tdImage1, err := ioutil.ReadFile("../testdata/image1.jpg")
	assert.NoError(t, err)

	ng := Engine{}
	imfo, err := ng.GetImageInfo(tdImage1)
	assert.NoError(t, err)

	assert.Equal(t, imfo.Width, 1600)
	assert.Equal(t, imfo.Height, 1200)
	assert.True(t, float64(int(imfo.AspectRatio*1000))/1000 == 1.333)
	assert.True(t, imfo.ContentLength == 451317)
}

func TestIssue8GIFResize(t *testing.T) {
	var sz *imgry.Sizing
	var img imgry.Image
	var err error

	ng := Engine{}

	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	assert.Equal(t, 131, img.Width())
	assert.Equal(t, 133, img.Height())

	origSize := len(img.Data())
	assert.Equal(t, 393324, origSize)

	img.Release()

	// Resizing to 750, which is slightly smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=750x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 750, img.Width())
	assert.Equal(t, 422, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.700.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 500, which is smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=500x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 500, img.Width())
	assert.Equal(t, 282, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.500.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 900, which is larger.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=900x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 900, img.Width())
	assert.Equal(t, 507, img.Height())

	assert.True(t, len(img.Data()) > origSize, fmt.Sprintf("Expecting %d > %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.900.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 200, which is smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=200x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 200, img.Width())
	assert.Equal(t, 113, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.200.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 150, which is smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=150x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 150, img.Width())
	assert.Equal(t, 84, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.150.gif")
	assert.NoError(t, err)

	img.Release()
}
<commit_msg>Disabling test temporarily.<commit_after>package imagick

import (
	"fmt"
	"io/ioutil"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/pressly/imgry"
)

func TestLoadBlob(t *testing.T) {
	tdImage1, err := ioutil.ReadFile("../testdata/image1.jpg")
	assert.NoError(t, err)

	ng := Engine{}
	im, err := ng.LoadBlob(tdImage1)
	assert.NoError(t, err)
	defer im.Release()

	sz, _ := imgry.NewSizingFromQuery("size=800x")
	err = im.SizeIt(sz)
	assert.NoError(t, err)

	im2Path := "/tmp/image1.jpg"
	im.WriteToFile(im2Path)

	im2, err := ng.LoadFile(im2Path)
	assert.NoError(t, err)

	assert.True(t, im2.Width() == 800)
	assert.Equal(t, "jpg", im2.Format())

	err = im2.SetFormat("png")
	assert.NoError(t, err)
}

func TestGetImageInfo(t *testing.T) {
	tdImage1, err := ioutil.ReadFile("../testdata/image1.jpg")
	assert.NoError(t, err)

	ng := Engine{}
	imfo, err := ng.GetImageInfo(tdImage1)
	assert.NoError(t, err)

	assert.Equal(t, imfo.Width, 1600)
	assert.Equal(t, imfo.Height, 1200)
	assert.True(t, float64(int(imfo.AspectRatio*1000))/1000 == 1.333)
	assert.True(t, imfo.ContentLength == 451317)
}

func TestIssue8GIFResize(t *testing.T) {
	var sz *imgry.Sizing
	var img imgry.Image
	var err error

	ng := Engine{}

	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	assert.Equal(t, 131, img.Width())
	assert.Equal(t, 133, img.Height())

	origSize := len(img.Data())
	assert.Equal(t, 393324, origSize)

	img.Release()

	// Resizing to 750, which is slightly smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=750x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 750, img.Width())
	assert.Equal(t, 422, img.Height())

	// We should be able to expect this someday, but for now it seems like the
	// number of colors after resizing affect the size of the file.
	//
	// See http://www.imagemagick.org/discourse-server/viewtopic.php?t=22505#p93859
	//assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.700.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 500, which is smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=500x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 500, img.Width())
	assert.Equal(t, 282, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.500.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 900, which is larger.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=900x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 900, img.Width())
	assert.Equal(t, 507, img.Height())

	assert.True(t, len(img.Data()) > origSize, fmt.Sprintf("Expecting %d > %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.900.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 200, which is smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=200x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 200, img.Width())
	assert.Equal(t, 113, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.200.gif")
	assert.NoError(t, err)

	img.Release()

	// Resizing to 150, which is smaller.
	img, err = ng.LoadFile("../testdata/issue-8.gif")
	assert.NoError(t, err)

	sz, _ = imgry.NewSizingFromQuery("size=150x")
	err = img.SizeIt(sz)
	assert.NoError(t, err)

	assert.Equal(t, 150, img.Width())
	assert.Equal(t, 84, img.Height())

	assert.True(t, len(img.Data()) < origSize, fmt.Sprintf("Expecting %d < %d.", len(img.Data()), origSize))

	err = img.WriteToFile("../testdata/issue-8.150.gif")
	assert.NoError(t, err)

	img.Release()
}
<|endoftext|>


<commit_before>package marketTypes

import (
	"net/http"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/antihax/goesi"
)

var esiClient goesi.APIClient
var esiSemaphore = make(chan struct{}, 500)
var typeIDs []int64

// Initialize initializes the market type updates
func Initialize() {
	httpClient := &http.Client{
		Timeout: time.Duration(time.Second * 10),
	}
	esiClient = *goesi.NewAPIClient(httpClient, "Element43/market-streamer (element-43.com)")

	updateTypes()
	go scheduleTypeUpdate()
}

// GetMarketTypes returns all typeIDs with a market
func GetMarketTypes() []int64 {
	return typeIDs
}

// Keep ticking in own goroutine and spawn worker tasks.
func scheduleTypeUpdate() {
	ticker := time.NewTicker(30 * time.Minute)
	defer ticker.Stop()
	for {
		<-ticker.C
		go updateTypes()
	}
}

// Update type list
func updateTypes() {
	logrus.Debug("Updating market types.")

	types, err := getMarketTypes()
	if err != nil {
		logrus.WithError(err).Error("Failed to get market types!")
	} else {
		typeIDs = types
	}

	logrus.Debug("Market type update done.")
}

// Get all types on market
func getMarketTypes() ([]int64, error) {
	typeIDs, err := getTypeIDs()
	if err != nil {
		return nil, err
	}

	marketTypes := make(chan int64)
	nonMarketTypes := make(chan int64)
	failure := make(chan error)

	typesLeft := len(typeIDs)

	for _, id := range typeIDs {
		go checkIfMarketTypeAsyncRetry(id, marketTypes, nonMarketTypes, failure)
	}

	var marketTypeIDs []int64

	for typesLeft > 0 {
		select {
		case typeID := <-marketTypes:
			marketTypeIDs = append(marketTypeIDs, typeID)
		case <-nonMarketTypes:
		case err := <-failure:
			logrus.Warnf("Error fetching type from ESI: %s", err.Error())
		}

		typesLeft--
	}

	return marketTypeIDs, nil
}

// Get all typeIDs from ESI
// TODO: move to static-data RPC
func getTypeIDs() ([]int32, error) {
	var typeIDs []int32
	params := make(map[string]interface{})
	params["page"] = int32(1)

	typeResult, _, err := esiClient.V1.UniverseApi.GetUniverseTypes(params)
	if err != nil {
		return nil, err
	}

	typeIDs = append(typeIDs, typeResult...)

	for len(typeResult) > 0 {
		params["page"] = params["page"].(int32) + 1
		typeResult, _, err = esiClient.V1.UniverseApi.GetUniverseTypes(params)
		if err != nil {
			return nil, err
		}

		typeIDs = append(typeIDs, typeResult...)
	}

	return typeIDs, nil
}

// Async check if market type, retry 3 times
func checkIfMarketTypeAsyncRetry(typeID int32, marketTypes chan int64, nonMarketTypes chan int64, failure chan error) {
	var isMarketType bool
	var err error
	retries := 3

	for retries > 0 {
		isMarketType, err = checkIfMarketType(typeID)
		if err != nil {
			retries--
		} else {
			err = nil
			retries = 0
		}
	}

	if err != nil {
		failure <- err
		return
	}

	if isMarketType {
		marketTypes <- int64(typeID)
		return
	}

	nonMarketTypes <- int64(typeID)
}

// Check if type is market type
func checkIfMarketType(typeID int32) (bool, error) {
	esiSemaphore <- struct{}{}
	typeInfo, _, err := esiClient.V3.UniverseApi.GetUniverseTypesTypeId(typeID, nil)
	<-esiSemaphore
	if err != nil {
		return false, err
	}

	// If it is published and has a market group it is a market type!
	if typeInfo.Published && (typeInfo.MarketGroupId != 0) {
		return true, nil
	}

	return false, nil
}
<commit_msg>Fetch types every 6 hours instead of 30 minutes<commit_after>package marketTypes

import (
	"net/http"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/antihax/goesi"
)

var esiClient goesi.APIClient
var esiSemaphore = make(chan struct{}, 500)
var typeIDs []int64

// Initialize initializes the market type updates
func Initialize() {
	httpClient := &http.Client{
		Timeout: time.Duration(time.Second * 10),
	}
	esiClient = *goesi.NewAPIClient(httpClient, "Element43/market-streamer (element-43.com)")

	updateTypes()
	go scheduleTypeUpdate()
}

// GetMarketTypes returns all typeIDs with a market
func GetMarketTypes() []int64 {
	return typeIDs
}

// Keep ticking in own goroutine and spawn worker tasks.
func scheduleTypeUpdate() {
	ticker := time.NewTicker(6 * time.Hour)
	defer ticker.Stop()
	for {
		<-ticker.C
		go updateTypes()
	}
}

// Update type list
func updateTypes() {
	logrus.Debug("Updating market types.")

	types, err := getMarketTypes()
	if err != nil {
		logrus.WithError(err).Error("Failed to get market types!")
	} else {
		typeIDs = types
	}

	logrus.Debug("Market type update done.")
}

// Get all types on market
func getMarketTypes() ([]int64, error) {
	typeIDs, err := getTypeIDs()
	if err != nil {
		return nil, err
	}

	marketTypes := make(chan int64)
	nonMarketTypes := make(chan int64)
	failure := make(chan error)

	typesLeft := len(typeIDs)

	for _, id := range typeIDs {
		go checkIfMarketTypeAsyncRetry(id, marketTypes, nonMarketTypes, failure)
	}

	var marketTypeIDs []int64

	for typesLeft > 0 {
		select {
		case typeID := <-marketTypes:
			marketTypeIDs = append(marketTypeIDs, typeID)
		case <-nonMarketTypes:
		case err := <-failure:
			logrus.Warnf("Error fetching type from ESI: %s", err.Error())
		}

		typesLeft--
	}

	return marketTypeIDs, nil
}

// Get all typeIDs from ESI
// TODO: move to static-data RPC
func getTypeIDs() ([]int32, error) {
	var typeIDs []int32
	params := make(map[string]interface{})
	params["page"] = int32(1)

	typeResult, _, err := esiClient.V1.UniverseApi.GetUniverseTypes(params)
	if err != nil {
		return nil, err
	}

	typeIDs = append(typeIDs, typeResult...)

	for len(typeResult) > 0 {
		params["page"] = params["page"].(int32) + 1
		typeResult, _, err = esiClient.V1.UniverseApi.GetUniverseTypes(params)
		if err != nil {
			return nil, err
		}

		typeIDs = append(typeIDs, typeResult...)
	}

	return typeIDs, nil
}

// Async check if market type, retry 3 times
func checkIfMarketTypeAsyncRetry(typeID int32, marketTypes chan int64, nonMarketTypes chan int64, failure chan error) {
	var isMarketType bool
	var err error
	retries := 3

	for retries > 0 {
		isMarketType, err = checkIfMarketType(typeID)
		if err != nil {
			retries--
		} else {
			err = nil
			retries = 0
		}
	}

	if err != nil {
		failure <- err
		return
	}

	if isMarketType {
		marketTypes <- int64(typeID)
		return
	}

	nonMarketTypes <- int64(typeID)
}

// Check if type is market type
func checkIfMarketType(typeID int32) (bool, error) {
	esiSemaphore <- struct{}{}
	typeInfo, _, err := esiClient.V3.UniverseApi.GetUniverseTypesTypeId(typeID, nil)
	<-esiSemaphore
	if err != nil {
		return false, err
	}

	// If it is published and has a market group it is a market type!
	if typeInfo.Published && (typeInfo.MarketGroupId != 0) {
		return true, nil
	}

	return false, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"gopkg.in/yaml.v2"
)

var removeKeys = []string{}

func main() {
	loadVariant()

	if len(flag.Args()) < 1 {
		fmt.Println("Usage: filter-crd <CRD YAML file>")
		return
	}

	f, err := os.Open(flag.Args()[0])
	if err != nil {
		log.Fatal("Error opening file", err)
	}

	decoder := yaml.NewDecoder(f)
	var d map[interface{}]interface{}
	output := []string{}

	for decoder.Decode(&d) == nil {

		if len(d) == 0 {
			continue
		}

		checkChain(d, []string{})

		fileOut, err := yaml.Marshal(d)
		if err != nil {
			log.Fatal("Error marshaling output", err)
		}

		output = append(output, string(fileOut))

	}

	fmt.Println(strings.Join(output, "---\n"))
}

func checkChain(d map[interface{}]interface{}, chain []string) {
	for k, v := range d {
		if key, ok := k.(string); ok {
			chain = append(chain, key)

			for _, removeKey := range removeKeys {
				if strings.Join(chain, "/") == removeKey {
					delete(d, key)
				}
			}

			if value, ok := v.(map[interface{}]interface{}); ok {
				checkChain(value, chain)
			}
			chain = chain[:len(chain)-1] // we're done with this key, remove it from the chain
		}
	}
}

func loadVariant() {
	variant := ""
	flag.StringVar(&variant, "variant", "", "variant of remove rules")
	flag.Parse()

	if variant == "cert-manager-openshift" {
		// These are the keys that the script will remove for OpenShift compatibility
		removeKeys = []string{
			"spec/preserveUnknownFields",
			"spec/validation/openAPIV3Schema/type",
		}
	}
}
<commit_msg>Improve errors<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"gopkg.in/yaml.v2"
)

var removeKeys = []string{}

func main() {
	loadVariant()

	if len(flag.Args()) < 1 {
		log.Fatal("Usage: filter-crd <CRD YAML file>")
	}

	f, err := os.Open(flag.Args()[0])
	if err != nil {
		log.Fatal("Error opening file: ", err)
	}

	decoder := yaml.NewDecoder(f)
	var d map[interface{}]interface{}
	output := []string{}

	for decoder.Decode(&d) == nil {

		if len(d) == 0 {
			continue
		}

		checkChain(d, []string{})

		fileOut, err := yaml.Marshal(d)
		if err != nil {
			log.Fatal("Error marshaling output: ", err)
		}

		output = append(output, string(fileOut))

	}

	fmt.Println(strings.Join(output, "---\n"))
}

func checkChain(d map[interface{}]interface{}, chain []string) {
	for k, v := range d {
		if key, ok := k.(string); ok {
			chain = append(chain, key)

			for _, removeKey := range removeKeys {
				if strings.Join(chain, "/") == removeKey {
					delete(d, key)
				}
			}

			if value, ok := v.(map[interface{}]interface{}); ok {
				checkChain(value, chain)
			}
			chain = chain[:len(chain)-1] // we're done with this key, remove it from the chain
		}
	}
}

func loadVariant() {
	variant := ""
	flag.StringVar(&variant, "variant", "", "variant of remove rules")
	flag.Parse()

	if variant == "cert-manager-openshift" {
		// These are the keys that the script will remove for OpenShift compatibility
		removeKeys = []string{
			"spec/preserveUnknownFields",
			"spec/validation/openAPIV3Schema/type",
		}
	}
}
<|endoftext|>


<commit_before>package autorest

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"time"
)

// Sender is the interface that wraps the Do method to send HTTP requests.
//
// The standard http.Client conforms to this interface.
type Sender interface {
	Do(*http.Request) (*http.Response, error)
}

// SenderFunc is a method that implements the Sender interface.
type SenderFunc func(*http.Request) (*http.Response, error)

// Do implements the Sender interface on SenderFunc.
func (sf SenderFunc) Do(r *http.Request) (*http.Response, error) {
	return sf(r)
}

// SendDecorator takes and possibily decorates, by wrapping, a Sender. Decorators may affect the
// http.Request and pass it along or, first, pass the http.Request along then react to the
// http.Response result.
type SendDecorator func(Sender) Sender

// CreateSender creates, decorates, and returns, as a Sender, the default http.Client.
func CreateSender(decorators ...SendDecorator) Sender {
	return DecorateSender(&http.Client{}, decorators...)
}

// DecorateSender accepts a Sender and a, possibly empty, set of SendDecorators, which is applies to
// the Sender. Decorators are applied in the order received, but their affect upon the request
// depends on whether they are a pre-decorator (change the http.Request and then pass it along) or a
// post-decorator (pass the http.Request along and react to the results in http.Response).
func DecorateSender(s Sender, decorators ...SendDecorator) Sender {
	for _, decorate := range decorators {
		s = decorate(s)
	}
	return s
}

// Send sends, by means of the default http.Client, the passed http.Request, returning the
// http.Response and possible error. It also accepts a, possibly empty, set of SendDecorators which
// it will apply the http.Client before invoking the Do method.
//
// Send is a convenience method and not recommended for production. Advanced users should use
// SendWithSender, passing and sharing their own Sender (e.g., instance of http.Client).
//
// Send will not poll or retry requests.
func Send(r *http.Request, decorators ...SendDecorator) (*http.Response, error) {
	return SendWithSender(&http.Client{}, r, decorators...)
}

// SendWithSender sends the passed http.Request, through the provided Sender, returning the
// http.Response and possible error. It also accepts a, possibly empty, set of SendDecorators which
// it will apply the http.Client before invoking the Do method.
//
// SendWithSender will not poll or retry requests.
func SendWithSender(s Sender, r *http.Request, decorators ...SendDecorator) (*http.Response, error) {
	return DecorateSender(s, decorators...).Do(r)
}

// AfterDelay returns a SendDecorator that delays for the passed time.Duration before
// invoking the Sender. The delay may be terminated by closing the optional channel on the
// http.Request. If canceled, no further Senders are invoked.
func AfterDelay(d time.Duration) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			if !DelayForBackoff(d, 1, r.Cancel) {
				return nil, fmt.Errorf("autorest: AfterDelay canceled before full delay")
			}
			return s.Do(r)
		})
	}
}

// AsIs returns a SendDecorator that invokes the passed Sender without modifying the http.Request.
func AsIs() SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			return s.Do(r)
		})
	}
}

// DoCloseIfError returns a SendDecorator that first invokes the passed Sender after which
// it closes the response if the passed Sender returns an error and the response body exists.
func DoCloseIfError() SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			resp, err := s.Do(r)
			if err != nil {
				Respond(resp, ByClosing())
			}
			return resp, err
		})
	}
}

// DoErrorIfStatusCode returns a SendDecorator that emits an error if the response StatusCode is
// among the set passed. Since these are artificial errors, the response body may still require
// closing.
func DoErrorIfStatusCode(codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			resp, err := s.Do(r)
			if err == nil && ResponseHasStatusCode(resp, codes...) {
				err = NewErrorWithResponse("autorest", "DoErrorIfStatusCode", resp, "%v %v failed with %s",
					resp.Request.Method,
					resp.Request.URL,
					resp.Status)
			}
			return resp, err
		})
	}
}

// DoErrorUnlessStatusCode returns a SendDecorator that emits an error unless the response
// StatusCode is among the set passed. Since these are artificial errors, the response body
// may still require closing.
func DoErrorUnlessStatusCode(codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			resp, err := s.Do(r)
			if err == nil && !ResponseHasStatusCode(resp, codes...) {
				err = NewErrorWithResponse("autorest", "DoErrorUnlessStatusCode", resp, "%v %v failed with %s",
					resp.Request.Method,
					resp.Request.URL,
					resp.Status)
			}
			return resp, err
		})
	}
}

// DoPollForStatusCodes returns a SendDecorator that polls if the http.Response contains one of the
// passed status codes. It expects the http.Response to contain a Location header providing the
// URL at which to poll (using GET) and will poll until the time passed is equal to or greater than
// the supplied duration. It will delay between requests for the duration specified in the
// RetryAfter header or, if the header is absent, the passed delay. Polling may be canceled by
// closing the optional channel on the http.Request.
func DoPollForStatusCodes(duration time.Duration, delay time.Duration, codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			resp, err = s.Do(r)

			if err == nil && ResponseHasStatusCode(resp, codes...) {
				r, err = NewPollingRequest(resp, r.Cancel)

				for err == nil && ResponseHasStatusCode(resp, codes...) {
					Respond(resp,
						ByClosing())
					resp, err = SendWithSender(s, r,
						AfterDelay(GetRetryAfter(resp, delay)))
				}
			}

			return resp, err
		})
	}
}

// DoRetryForAttempts returns a SendDecorator that retries a failed request for up to the specified
// number of attempts, exponentially backing off between requests using the supplied backoff
// time.Duration (which may be zero). Retrying may be canceled by closing the optional channel on
// the http.Request.
func DoRetryForAttempts(attempts int, backoff time.Duration) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			for attempt := 0; attempt < attempts; attempt++ {
				resp, err = s.Do(r)
				if err == nil {
					return resp, err
				}
				DelayForBackoff(backoff, attempt, r.Cancel)
			}
			return resp, err
		})
	}
}

// DoRetryForStatusCodes returns a SendDecorator that retries for specified statusCodes for up to the specified
// number of attempts, exponentially backing off between requests using the supplied backoff
// time.Duration (which may be zero). Retrying may be canceled by closing the optional channel on
// the http.Request.
func DoRetryForStatusCodes(attempts int, backoff time.Duration, codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			b := []byte{}
			if r.Body != nil {
				b, err = ioutil.ReadAll(r.Body)
				if err != nil {
					return resp, err
				}
			}

			// Increment to add the first call (attempts denotes number of retries)
			attempts++
			for attempt := 0; attempt < attempts; attempt++ {
				r.Body = ioutil.NopCloser(bytes.NewBuffer(b))
				resp, err = s.Do(r)
				if err != nil || !ResponseHasStatusCode(resp, codes...) {
					return resp, err
				}
				DelayForBackoff(backoff, attempt, r.Cancel)
			}
			return resp, err
		})
	}
}

// DoRetryForDuration returns a SendDecorator that retries the request until the total time is equal
// to or greater than the specified duration, exponentially backing off between requests using the
// supplied backoff time.Duration (which may be zero). Retrying may be canceled by closing the
// optional channel on the http.Request.
func DoRetryForDuration(d time.Duration, backoff time.Duration) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			end := time.Now().Add(d)
			for attempt := 0; time.Now().Before(end); attempt++ {
				resp, err = s.Do(r)
				if err == nil {
					return resp, err
				}
				DelayForBackoff(backoff, attempt, r.Cancel)
			}
			return resp, err
		})
	}
}

// WithLogging returns a SendDecorator that implements simple before and after logging of the
// request.
func WithLogging(logger *log.Logger) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			logger.Printf("Sending %s %s", r.Method, r.URL)
			resp, err := s.Do(r)
			if err != nil {
				logger.Printf("%s %s received error '%v'", r.Method, r.URL, err)
			} else {
				logger.Printf("%s %s received %s", r.Method, r.URL, resp.Status)
			}
			return resp, err
		})
	}
}

// DelayForBackoff invokes time.After for the supplied backoff duration raised to the power of
// passed attempt (i.e., an exponential backoff delay). Backoff duration is in seconds and can set
// to zero for no delay. The delay may be canceled by closing the passed channel. If terminated early,
// returns false.
func DelayForBackoff(backoff time.Duration, attempt int, cancel <-chan struct{}) bool {
	select {
	case <-time.After(time.Duration(backoff.Seconds()*math.Pow(2, float64(attempt))) * time.Second):
		return true
	case <-cancel:
		return false
	}
}
<commit_msg>Pass atempt 0 to DelayForBackoff in AfterDelay<commit_after>package autorest

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"time"
)

// Sender is the interface that wraps the Do method to send HTTP requests.
//
// The standard http.Client conforms to this interface.
type Sender interface {
	Do(*http.Request) (*http.Response, error)
}

// SenderFunc is a method that implements the Sender interface.
type SenderFunc func(*http.Request) (*http.Response, error)

// Do implements the Sender interface on SenderFunc.
func (sf SenderFunc) Do(r *http.Request) (*http.Response, error) {
	return sf(r)
}

// SendDecorator takes and possibily decorates, by wrapping, a Sender. Decorators may affect the
// http.Request and pass it along or, first, pass the http.Request along then react to the
// http.Response result.
type SendDecorator func(Sender) Sender

// CreateSender creates, decorates, and returns, as a Sender, the default http.Client.
func CreateSender(decorators ...SendDecorator) Sender {
	return DecorateSender(&http.Client{}, decorators...)
}

// DecorateSender accepts a Sender and a, possibly empty, set of SendDecorators, which is applies to
// the Sender. Decorators are applied in the order received, but their affect upon the request
// depends on whether they are a pre-decorator (change the http.Request and then pass it along) or a
// post-decorator (pass the http.Request along and react to the results in http.Response).
func DecorateSender(s Sender, decorators ...SendDecorator) Sender {
	for _, decorate := range decorators {
		s = decorate(s)
	}
	return s
}

// Send sends, by means of the default http.Client, the passed http.Request, returning the
// http.Response and possible error. It also accepts a, possibly empty, set of SendDecorators which
// it will apply the http.Client before invoking the Do method.
//
// Send is a convenience method and not recommended for production. Advanced users should use
// SendWithSender, passing and sharing their own Sender (e.g., instance of http.Client).
//
// Send will not poll or retry requests.
func Send(r *http.Request, decorators ...SendDecorator) (*http.Response, error) {
	return SendWithSender(&http.Client{}, r, decorators...)
}

// SendWithSender sends the passed http.Request, through the provided Sender, returning the
// http.Response and possible error. It also accepts a, possibly empty, set of SendDecorators which
// it will apply the http.Client before invoking the Do method.
//
// SendWithSender will not poll or retry requests.
func SendWithSender(s Sender, r *http.Request, decorators ...SendDecorator) (*http.Response, error) {
	return DecorateSender(s, decorators...).Do(r)
}

// AfterDelay returns a SendDecorator that delays for the passed time.Duration before
// invoking the Sender. The delay may be terminated by closing the optional channel on the
// http.Request. If canceled, no further Senders are invoked.
func AfterDelay(d time.Duration) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			if !DelayForBackoff(d, 0, r.Cancel) {
				return nil, fmt.Errorf("autorest: AfterDelay canceled before full delay")
			}
			return s.Do(r)
		})
	}
}

// AsIs returns a SendDecorator that invokes the passed Sender without modifying the http.Request.
func AsIs() SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			return s.Do(r)
		})
	}
}

// DoCloseIfError returns a SendDecorator that first invokes the passed Sender after which
// it closes the response if the passed Sender returns an error and the response body exists.
func DoCloseIfError() SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			resp, err := s.Do(r)
			if err != nil {
				Respond(resp, ByClosing())
			}
			return resp, err
		})
	}
}

// DoErrorIfStatusCode returns a SendDecorator that emits an error if the response StatusCode is
// among the set passed. Since these are artificial errors, the response body may still require
// closing.
func DoErrorIfStatusCode(codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			resp, err := s.Do(r)
			if err == nil && ResponseHasStatusCode(resp, codes...) {
				err = NewErrorWithResponse("autorest", "DoErrorIfStatusCode", resp, "%v %v failed with %s",
					resp.Request.Method,
					resp.Request.URL,
					resp.Status)
			}
			return resp, err
		})
	}
}

// DoErrorUnlessStatusCode returns a SendDecorator that emits an error unless the response
// StatusCode is among the set passed. Since these are artificial errors, the response body
// may still require closing.
func DoErrorUnlessStatusCode(codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			resp, err := s.Do(r)
			if err == nil && !ResponseHasStatusCode(resp, codes...) {
				err = NewErrorWithResponse("autorest", "DoErrorUnlessStatusCode", resp, "%v %v failed with %s",
					resp.Request.Method,
					resp.Request.URL,
					resp.Status)
			}
			return resp, err
		})
	}
}

// DoPollForStatusCodes returns a SendDecorator that polls if the http.Response contains one of the
// passed status codes. It expects the http.Response to contain a Location header providing the
// URL at which to poll (using GET) and will poll until the time passed is equal to or greater than
// the supplied duration. It will delay between requests for the duration specified in the
// RetryAfter header or, if the header is absent, the passed delay. Polling may be canceled by
// closing the optional channel on the http.Request.
func DoPollForStatusCodes(duration time.Duration, delay time.Duration, codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			resp, err = s.Do(r)

			if err == nil && ResponseHasStatusCode(resp, codes...) {
				r, err = NewPollingRequest(resp, r.Cancel)

				for err == nil && ResponseHasStatusCode(resp, codes...) {
					Respond(resp,
						ByClosing())
					resp, err = SendWithSender(s, r,
						AfterDelay(GetRetryAfter(resp, delay)))
				}
			}

			return resp, err
		})
	}
}

// DoRetryForAttempts returns a SendDecorator that retries a failed request for up to the specified
// number of attempts, exponentially backing off between requests using the supplied backoff
// time.Duration (which may be zero). Retrying may be canceled by closing the optional channel on
// the http.Request.
func DoRetryForAttempts(attempts int, backoff time.Duration) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			for attempt := 0; attempt < attempts; attempt++ {
				resp, err = s.Do(r)
				if err == nil {
					return resp, err
				}
				DelayForBackoff(backoff, attempt, r.Cancel)
			}
			return resp, err
		})
	}
}

// DoRetryForStatusCodes returns a SendDecorator that retries for specified statusCodes for up to the specified
// number of attempts, exponentially backing off between requests using the supplied backoff
// time.Duration (which may be zero). Retrying may be canceled by closing the optional channel on
// the http.Request.
func DoRetryForStatusCodes(attempts int, backoff time.Duration, codes ...int) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			b := []byte{}
			if r.Body != nil {
				b, err = ioutil.ReadAll(r.Body)
				if err != nil {
					return resp, err
				}
			}

			// Increment to add the first call (attempts denotes number of retries)
			attempts++
			for attempt := 0; attempt < attempts; attempt++ {
				r.Body = ioutil.NopCloser(bytes.NewBuffer(b))
				resp, err = s.Do(r)
				if err != nil || !ResponseHasStatusCode(resp, codes...) {
					return resp, err
				}
				DelayForBackoff(backoff, attempt, r.Cancel)
			}
			return resp, err
		})
	}
}

// DoRetryForDuration returns a SendDecorator that retries the request until the total time is equal
// to or greater than the specified duration, exponentially backing off between requests using the
// supplied backoff time.Duration (which may be zero). Retrying may be canceled by closing the
// optional channel on the http.Request.
func DoRetryForDuration(d time.Duration, backoff time.Duration) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
			end := time.Now().Add(d)
			for attempt := 0; time.Now().Before(end); attempt++ {
				resp, err = s.Do(r)
				if err == nil {
					return resp, err
				}
				DelayForBackoff(backoff, attempt, r.Cancel)
			}
			return resp, err
		})
	}
}

// WithLogging returns a SendDecorator that implements simple before and after logging of the
// request.
func WithLogging(logger *log.Logger) SendDecorator {
	return func(s Sender) Sender {
		return SenderFunc(func(r *http.Request) (*http.Response, error) {
			logger.Printf("Sending %s %s", r.Method, r.URL)
			resp, err := s.Do(r)
			if err != nil {
				logger.Printf("%s %s received error '%v'", r.Method, r.URL, err)
			} else {
				logger.Printf("%s %s received %s", r.Method, r.URL, resp.Status)
			}
			return resp, err
		})
	}
}

// DelayForBackoff invokes time.After for the supplied backoff duration raised to the power of
// passed attempt (i.e., an exponential backoff delay). Backoff duration is in seconds and can set
// to zero for no delay. The delay may be canceled by closing the passed channel. If terminated early,
// returns false.
func DelayForBackoff(backoff time.Duration, attempt int, cancel <-chan struct{}) bool {
	select {
	case <-time.After(time.Duration(backoff.Seconds()*math.Pow(2, float64(attempt))) * time.Second):
		return true
	case <-cancel:
		return false
	}
}
<|endoftext|>


<commit_before>package entity

import (
	"sync"

	. "chunkymonkey/types"
)

// TODO EntityManager should be a service in its own right, able to hand out
// blocks of IDs and running its own goroutine (potentially shardable by
// entityId if necessary). Right now taking the easy option of using a simple
// lock.
type EntityManager struct {
	nextEntityId EntityId
	entities     map[EntityId]bool
	lock         sync.Mutex
}

func (mgr *EntityManager) Init() {
	mgr.lock.Lock()
	defer mgr.lock.Unlock()

	mgr.nextEntityId = 0
	mgr.entities = make(map[EntityId]bool)
}

func (mgr *EntityManager) createEntityId() EntityId {
	// Search for next free ID
	entityId := mgr.nextEntityId
	_, exists := mgr.entities[entityId]
	for exists {
		entityId++
		if entityId == mgr.nextEntityId {
			// TODO Better handling of this? It shouldn't happen, realistically - but
			// neither should it explode.
			panic("EntityId space exhausted")
		}
		_, exists = mgr.entities[entityId]
	}
	mgr.nextEntityId = entityId + 1

	return entityId
}

// AddEntity adds an entity to the manager, and assigns it a world-unique
// EntityId.
// NewEntity creates a world-unique entityId in the manager and returns it.
func (mgr *EntityManager) NewEntity() EntityId {
	mgr.lock.Lock()
	defer mgr.lock.Unlock()

	entityId := mgr.createEntityId()
	mgr.entities[entityId] = true
	return entityId
}

// RemoveEntity removes an entity from the manager.
func (mgr *EntityManager) RemoveEntityById(entityId EntityId) {
	mgr.lock.Lock()
	defer mgr.lock.Unlock()

	mgr.entities[entityId] = false, false
}
<commit_msg>Removing a TODO comment that's thinking too far ahead.<commit_after>package entity

import (
	"sync"

	. "chunkymonkey/types"
)

type EntityManager struct {
	nextEntityId EntityId
	entities     map[EntityId]bool
	lock         sync.Mutex
}

func (mgr *EntityManager) Init() {
	mgr.lock.Lock()
	defer mgr.lock.Unlock()

	mgr.nextEntityId = 0
	mgr.entities = make(map[EntityId]bool)
}

func (mgr *EntityManager) createEntityId() EntityId {
	// Search for next free ID
	entityId := mgr.nextEntityId
	_, exists := mgr.entities[entityId]
	for exists {
		entityId++
		if entityId == mgr.nextEntityId {
			// TODO Better handling of this? It shouldn't happen, realistically - but
			// neither should it explode.
			panic("EntityId space exhausted")
		}
		_, exists = mgr.entities[entityId]
	}
	mgr.nextEntityId = entityId + 1

	return entityId
}

// AddEntity adds an entity to the manager, and assigns it a world-unique
// EntityId.
// NewEntity creates a world-unique entityId in the manager and returns it.
func (mgr *EntityManager) NewEntity() EntityId {
	mgr.lock.Lock()
	defer mgr.lock.Unlock()

	entityId := mgr.createEntityId()
	mgr.entities[entityId] = true
	return entityId
}

// RemoveEntity removes an entity from the manager.
func (mgr *EntityManager) RemoveEntityById(entityId EntityId) {
	mgr.lock.Lock()
	defer mgr.lock.Unlock()

	mgr.entities[entityId] = false, false
}
<|endoftext|>


<commit_before>package avltree

// go fmt . && go test . -v

// based on http://www.geeksforgeeks.org/avl-tree-set-1-insertion/

type Node struct {
	Key    int
	Left   *Node
	Right  *Node
	height int
}

func createNode(key int) *Node {
	n := &Node{}
	n.Key = key
	n.height = 1

	// initialise other properties
	// ...

	return n
}

func height(n *Node) int {
	if n == nil {
		return 0
	}
	return n.height
}

func max(a int, b int) int {
	if a > b {
		return a
	}
	return b
}

/*
       y                           x
      / \     rotateRight()       / \
     x  T3    – – – – – – >     T1  y
    / \       < - - - - - -        / \
   T1  T2      rotateLeft()      T2  T3
*/

func rotateLeft(x *Node) *Node {
	y := x.Right
	t2 := y.Left
	y.Left = x
	x.Right = t2
	x.height = max(height(x.Left), height(x.Right)) + 1
	y.height = max(height(y.Left), height(y.Right)) + 1

	// update node invariants
	// ...

	return y
}

func rotateRight(y *Node) *Node {
	x := y.Left
	t2 := x.Right
	x.Right = y
	y.Left = t2
	y.height = max(height(y.Left), height(y.Right)) + 1
	x.height = max(height(x.Left), height(x.Right)) + 1

	// update node invariants
	// ...

	return x
}

func balance(n *Node, key int) *Node {
	b := height(n.Left) - height(n.Right)

	if b > 1 && key < n.Left.Key {
		return rotateRight(n)
	}

	if b < -1 && key > n.Right.Key {
		return rotateLeft(n)
	}

	if b > 1 && key > n.Left.Key {
		n.Left = rotateLeft(n.Left)
		return rotateRight(n)
	}

	if b < -1 && key < n.Right.Key {
		n.Right = rotateRight(n.Right)
		return rotateLeft(n)
	}

	return n
}

func InsertOrIgnore(n *Node, key int) *Node {
	if n == nil {
		return createNode(key)
	}

	// update statistics on the way down
	// ...

	if key < n.Key {
		n.Left = Insert(n.Left, key)
	} else if key > n.Key {
		n.Right = Insert(n.Right, key)
	} else if key == n.Key {
		return n
	}

	n.height = max(height(n.Left), height(n.Right)) + 1

	return balance(n, key)
}

func InsertOrReplace(n *Node, key int) *Node {
	if n == nil {
		return createNode(key)
	}

	// update statistics on the way down
	// ...

	if key < n.Key {
		n.Left = Insert(n.Left, key)
	} else if key > n.Key {
		n.Right = Insert(n.Right, key)
	} else if key == n.Key {
		cn := createNode(key)
		cn.height = n.height
		cn.Left = n.Left
		cn.Right = n.Right

		// copy node statistics
		// ...

		return cn
	}

	n.height = max(height(n.Left), height(n.Right)) + 1

	return balance(n, key)
}

func Insert(n *Node, key int) *Node {
	if n == nil {
		return createNode(key)
	}

	// update statistics on the way down
	// ...

	if key < n.Key {
		n.Left = Insert(n.Left, key)
	} else {
		// duplicates go right
		n.Right = Insert(n.Right, key)
	}

	n.height = max(height(n.Left), height(n.Right)) + 1

	return balance(n, key)
}
<commit_msg>Update avltree.go<commit_after>package avltree

// go fmt . && go test . -v

// based on http://www.geeksforgeeks.org/avl-tree-set-1-insertion/

type Node struct {
	Key    int
	Left   *Node
	Right  *Node
	height int
}

func createNode(key int) *Node {
	n := &Node{}
	n.Key = key
	n.height = 1

	// initialise other properties
	// ...

	return n
}

func height(n *Node) int {
	if n == nil {
		return 0
	}
	return n.height
}

func max(a int, b int) int {
	if a > b {
		return a
	}
	return b
}

/*
       y                           x
      / \     rotateRight()       / \
     x  T3    – – – – – – >     T1  y
    / \       < - - - - - -        / \
   T1  T2      rotateLeft()      T2  T3
*/

func rotateLeft(x *Node) *Node {
	y := x.Right
	t2 := y.Left
	y.Left = x
	x.Right = t2
	x.height = max(height(x.Left), height(x.Right)) + 1
	y.height = max(height(y.Left), height(y.Right)) + 1

	// update node invariants
	// ...

	return y
}

func rotateRight(y *Node) *Node {
	x := y.Left
	t2 := x.Right
	x.Right = y
	y.Left = t2
	y.height = max(height(y.Left), height(y.Right)) + 1
	x.height = max(height(x.Left), height(x.Right)) + 1

	// update node invariants
	// ...

	return x
}

func balance(n *Node, key int) *Node {
	b := height(n.Left) - height(n.Right)

	if b > 1 && key < n.Left.Key {
		return rotateRight(n)
	}

	if b < -1 && key > n.Right.Key {
		return rotateLeft(n)
	}

	if b > 1 && key > n.Left.Key {
		n.Left = rotateLeft(n.Left)
		return rotateRight(n)
	}

	if b < -1 && key < n.Right.Key {
		n.Right = rotateRight(n.Right)
		return rotateLeft(n)
	}

	return n
}

func InsertOrIgnore(n *Node, key int) *Node {
	if n == nil {
		return createNode(key)
	}

	// update statistics on the way down
	// ...

	if key < n.Key {
		n.Left = InsertOrIgnore(n.Left, key)
	} else if key > n.Key {
		n.Right = InsertOrIgnore(n.Right, key)
	} else if key == n.Key {
		return n
	}

	n.height = max(height(n.Left), height(n.Right)) + 1

	return balance(n, key)
}

func InsertOrReplace(n *Node, key int) *Node {
	if n == nil {
		return createNode(key)
	}

	// update statistics on the way down
	// ...

	if key < n.Key {
		n.Left = InsertOrReplace(n.Left, key)
	} else if key > n.Key {
		n.Right = InsertOrReplace(n.Right, key)
	} else if key == n.Key {
		cn := createNode(key)
		cn.height = n.height
		cn.Left = n.Left
		cn.Right = n.Right

		// copy node statistics
		// ...

		return cn
	}

	n.height = max(height(n.Left), height(n.Right)) + 1

	return balance(n, key)
}

func Insert(n *Node, key int) *Node {
	if n == nil {
		return createNode(key)
	}

	// update statistics on the way down
	// ...

	if key < n.Key {
		n.Left = Insert(n.Left, key)
	} else {
		// duplicates go right
		n.Right = Insert(n.Right, key)
	}

	n.height = max(height(n.Left), height(n.Right)) + 1

	return balance(n, key)
}
<|endoftext|>


<commit_before>package integration

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"regexp"
	"strconv"
	"text/template"
	"time"

	yaml "gopkg.in/yaml.v2"

	"os"
	"os/exec"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/jmcvetta/guid"
)

var guidMaker = guid.SimpleGenerator()

const TARGET_REGION = "us-east-1"
const SUBNETID = "subnet-85f111b9"
const KEYNAME = "kismatic-integration-testing"
const SECURITYGROUPID = "sg-d1dc4dab"
const AMIUbuntu1604USEAST = "ami-29f96d3e"
const AMICentos7UsEast = "ami-6d1c2007"

var _ = Describe("Happy Path Installation Tests", func() {
	kisPath := CopyKismaticToTemp()

	BeforeSuite(func() {
		fmt.Println("Unpacking kismatic to", kisPath)
		c := exec.Command("tar", "-zxf", "../out/kismatic.tar.gz", "-C", kisPath)
		tarOut, tarErr := c.CombinedOutput()
		if tarErr != nil {
			log.Fatal("Error unpacking installer", string(tarOut), tarErr)
		}
		os.Chdir(kisPath)
	})

	AfterSuite(func() {
		//os.RemoveAll(kisPath)
	})

	Describe("Calling installer with no input", func() {
		It("should output help text", func() {
			c := exec.Command("./kismatic")
			helpbytes, helperr := c.Output()
			Expect(helperr).To(BeNil())
			helpText := string(helpbytes)
			Expect(helpText).To(ContainSubstring("Usage"))
		})
	})

	Describe("Calling installer with 'install plan'", func() {
		Context("and just hitting enter", func() {
			It("should result in the output of a well formed default plan file", func() {
				By("Outputing a file")
				c := exec.Command("./kismatic", "install", "plan")
				helpbytes, helperr := c.Output()
				Expect(helperr).To(BeNil())
				helpText := string(helpbytes)
				Expect(helpText).To(ContainSubstring("Generating installation plan file with 3 etcd nodes, 2 master nodes and 3 worker nodes"))
				Expect(FileExists("kismatic-cluster.yaml")).To(Equal(true))

				By("Outputing a file with valid YAML")
				yamlBytes, err := ioutil.ReadFile("kismatic-cluster.yaml")
				if err != nil {
					Fail("Could not read cluster file")
				}
				yamlBlob := string(yamlBytes)

				planFromYaml := ClusterPlan{}

				unmarshallErr := yaml.Unmarshal([]byte(yamlBlob), &planFromYaml)
				if unmarshallErr != nil {
					Fail("Could not unmarshall cluster yaml: %v")
				}
			})
		})
	})

	Describe("Calling installer with a plan targetting AWS", func() {
		// Context("Using a 1/1/1 Ubtunu 16.04 layout", func() {
		// 	It("should result in a working cluster", func() {
		// 		InstallKismatic(AMIUbuntu1604USEAST, "ubuntu")
		// 	})
		// })
		Context("Using a 1/1/1 CentOS 7 layout", func() {
			It("should result in a working cluster", func() {
				InstallKismatic(AMICentos7UsEast, "centos")
			})
		})
	})
})

func InstallKismatic(nodeType string, user string) {
	By("Building a template")
	template, err := template.New("planAWSOverlay").Parse(planAWSOverlay)
	FailIfError(err, "Couldn't parse template")

	By("Making infrastructure")
	etcdNode, etcErr := MakeETCDNode(nodeType)
	FailIfError(etcErr, "Error making etcd node")

	masterNode, masterErr := MakeMasterNode(nodeType)
	FailIfError(masterErr, "Error making master node")

	workerNode, workerErr := MakeWorkerNode(nodeType)
	FailIfError(workerErr, "Error making worker node")

	//defer TerminateInstances(etcdNode.Instanceid, masterNode.Instanceid, workerNode.Instanceid)
	descErr := WaitForInstanceToStart(&etcdNode, &masterNode, &workerNode)
	FailIfError(descErr, "Error waiting for nodes")
	log.Printf("Created etcd nodes: %v (%v), master nodes %v (%v), workerNodes %v (%v)",
		etcdNode.Instanceid, etcdNode.Publicip,
		masterNode.Instanceid, masterNode.Publicip,
		workerNode.Instanceid, workerNode.Publicip)

	By("Building a plan to set up an overlay network cluster on this hardware")
	nodes := PlanUbuntuAWS{
		Etcd:                []AWSNodeDeets{etcdNode},
		Master:              []AWSNodeDeets{masterNode},
		Worker:              []AWSNodeDeets{workerNode},
		MasterNodeFQDN:      masterNode.Hostname,
		MasterNodeShortName: masterNode.Hostname,
		User:                user,
	}
	f, fileErr := os.Create("kismatic-testing.yaml")
	FailIfError(fileErr, "Error waiting for nodes")
	defer f.Close()
	w := bufio.NewWriter(f)
	execErr := template.Execute(w, &nodes)
	FailIfError(execErr, "Error filling in plan template")
	w.Flush()

	By("Validing our plan")
	ver := exec.Command("./kismatic", "install", "validate", "-f", f.Name())
	verbytes, verErr := ver.CombinedOutput()
	verText := string(verbytes)

	FailIfError(verErr, "Error validating plan", verText)

	By("Punch it Chewie!")
	app := exec.Command("./kismatic", "install", "apply", "-f", f.Name())
	appbytes, appErr := app.CombinedOutput()
	appText := string(appbytes)

	FailIfError(appErr, "Error applying plan", appText)
}

func FailIfError(err error, message ...string) {
	if err != nil {
		log.Printf(message[0]+": %v\n%v", err, message[1:])
		Fail(message[0])
	}
}

func CopyKismaticToTemp() string {
	tmpDir := os.TempDir()
	randomness, randomErr := GenerateGUIDString()
	if randomErr != nil {
		log.Fatal("Error making a GUID: ", randomErr)
	}
	kisPath := tmpDir + "/kisint/" + randomness
	err := os.MkdirAll(kisPath, 0777)
	if err != nil {
		log.Fatal("Error making temp dir: ", err)
	}

	return kisPath
}

func GenerateGUIDString() (string, error) {
	randomness, randomErr := guidMaker.NextId()

	if randomErr != nil {
		return "", randomErr
	}

	return strconv.FormatInt(randomness, 16), nil
}

func AssertKismaticDirectory(kisPath string) {
	if FileExists(kisPath + "/kismatic") {
		log.Fatal("Installer unpacked but kismatic wasn't there")
	}
}

func FileExists(path string) bool {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}

func MakeETCDNode(nodeType string) (AWSNodeDeets, error) {
	return MakeAWSNode(nodeType, ec2.InstanceTypeT2Micro)
}

func MakeMasterNode(nodeType string) (AWSNodeDeets, error) {
	return MakeAWSNode(nodeType, ec2.InstanceTypeT2Micro)
}

func MakeWorkerNode(nodeType string) (AWSNodeDeets, error) {
	return MakeAWSNode(nodeType, ec2.InstanceTypeT2Medium)
}

func MakeAWSNode(ami string, instanceType string) (AWSNodeDeets, error) {
	svc := ec2.New(session.New(&aws.Config{Region: aws.String(TARGET_REGION)}))
	runResult, err := svc.RunInstances(&ec2.RunInstancesInput{
		ImageId:          aws.String(ami),
		InstanceType:     aws.String(instanceType),
		MinCount:         aws.Int64(1),
		MaxCount:         aws.Int64(1),
		SubnetId:         aws.String(SUBNETID),
		KeyName:          aws.String(KEYNAME),
		SecurityGroupIds: []*string{aws.String(SECURITYGROUPID)},
	})

	if err != nil {
		return AWSNodeDeets{}, err
	}

	re := regexp.MustCompile("[^.]*")
	hostname := re.FindString(*runResult.Instances[0].PrivateDnsName)

	deets := AWSNodeDeets{
		Instanceid: *runResult.Instances[0].InstanceId,
		Privateip:  *runResult.Instances[0].PrivateIpAddress,
		Hostname:   hostname,
	}

	_, errtag := svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{aws.String(deets.Instanceid)},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("ApprendaTeam"),
				Value: aws.String("Kismatic"),
			},
		},
	})
	if errtag != nil {
		return deets, errtag
	}

	return deets, nil
}

func TerminateInstances(instanceids ...string) {
	awsinstanceids := make([]*string, len(instanceids))
	for i, id := range instanceids {
		awsinstanceids[i] = aws.String(id)
	}
	sess, err := session.NewSession()

	if err != nil {
		log.Printf("failed to create session: %v", err)
		return
	}

	svc := ec2.New(sess, &aws.Config{Region: aws.String(TARGET_REGION)})

	params := &ec2.TerminateInstancesInput{
		InstanceIds: awsinstanceids,
	}
	resp, err := svc.TerminateInstances(params)

	if err != nil {
		log.Printf("Could not terminate: %v", resp)
		return
	}
}

func WaitForInstanceToStart(nodes ...*AWSNodeDeets) error {
	sess, err := session.NewSession()

	if err != nil {
		fmt.Println("failed to create session,", err)
		return err
	}

	fmt.Print("Waiting for nodes to come up")
	defer fmt.Println()

	svc := ec2.New(sess, &aws.Config{Region: aws.String(TARGET_REGION)})
	for _, deets := range nodes {
		deets.Publicip = ""

		for deets.Publicip == "" {
			fmt.Print(".")
			descResult, descErr := svc.DescribeInstances(&ec2.DescribeInstancesInput{
				InstanceIds: []*string{aws.String(deets.Instanceid)},
			})
			if descErr != nil {
				return descErr
			}

			if *descResult.Reservations[0].Instances[0].State.Name == ec2.InstanceStateNameRunning &&
				descResult.Reservations[0].Instances[0].PublicIpAddress != nil {
				deets.Publicip = *descResult.Reservations[0].Instances[0].PublicIpAddress
				BlockUntilSSHOpen(deets)
			} else {
				time.Sleep(1 * time.Second)
			}
		}
	}
	return nil
}

func BlockUntilSSHOpen(node *AWSNodeDeets) {
	conn, err := net.Dial("tcp", node.Publicip+":22")
	fmt.Print("?")
	if err != nil {
		time.Sleep(5 & time.Second)
		BlockUntilSSHOpen(node)
	} else {
		conn.Close()
		return
	}
}
<commit_msg>Forgot the code to clean up resources after a run.<commit_after>package integration

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"regexp"
	"strconv"
	"text/template"
	"time"

	yaml "gopkg.in/yaml.v2"

	"os"
	"os/exec"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/jmcvetta/guid"
)

var guidMaker = guid.SimpleGenerator()

const TARGET_REGION = "us-east-1"
const SUBNETID = "subnet-85f111b9"
const KEYNAME = "kismatic-integration-testing"
const SECURITYGROUPID = "sg-d1dc4dab"
const AMIUbuntu1604USEAST = "ami-29f96d3e"
const AMICentos7UsEast = "ami-6d1c2007"

var _ = Describe("Happy Path Installation Tests", func() {
	kisPath := CopyKismaticToTemp()

	BeforeSuite(func() {
		fmt.Println("Unpacking kismatic to", kisPath)
		c := exec.Command("tar", "-zxf", "../out/kismatic.tar.gz", "-C", kisPath)
		tarOut, tarErr := c.CombinedOutput()
		if tarErr != nil {
			log.Fatal("Error unpacking installer", string(tarOut), tarErr)
		}
		os.Chdir(kisPath)
	})

	AfterSuite(func() {
		os.RemoveAll(kisPath)
	})

	Describe("Calling installer with no input", func() {
		It("should output help text", func() {
			c := exec.Command("./kismatic")
			helpbytes, helperr := c.Output()
			Expect(helperr).To(BeNil())
			helpText := string(helpbytes)
			Expect(helpText).To(ContainSubstring("Usage"))
		})
	})

	Describe("Calling installer with 'install plan'", func() {
		Context("and just hitting enter", func() {
			It("should result in the output of a well formed default plan file", func() {
				By("Outputing a file")
				c := exec.Command("./kismatic", "install", "plan")
				helpbytes, helperr := c.Output()
				Expect(helperr).To(BeNil())
				helpText := string(helpbytes)
				Expect(helpText).To(ContainSubstring("Generating installation plan file with 3 etcd nodes, 2 master nodes and 3 worker nodes"))
				Expect(FileExists("kismatic-cluster.yaml")).To(Equal(true))

				By("Outputing a file with valid YAML")
				yamlBytes, err := ioutil.ReadFile("kismatic-cluster.yaml")
				if err != nil {
					Fail("Could not read cluster file")
				}
				yamlBlob := string(yamlBytes)

				planFromYaml := ClusterPlan{}

				unmarshallErr := yaml.Unmarshal([]byte(yamlBlob), &planFromYaml)
				if unmarshallErr != nil {
					Fail("Could not unmarshall cluster yaml: %v")
				}
			})
		})
	})

	Describe("Calling installer with a plan targetting AWS", func() {
		Context("Using a 1/1/1 Ubtunu 16.04 layout", func() {
			It("should result in a working cluster", func() {
				InstallKismatic(AMIUbuntu1604USEAST, "ubuntu")
			})
		})
		Context("Using a 1/1/1 CentOS 7 layout", func() {
			It("should result in a working cluster", func() {
				InstallKismatic(AMICentos7UsEast, "centos")
			})
		})
	})
})

func InstallKismatic(nodeType string, user string) {
	By("Building a template")
	template, err := template.New("planAWSOverlay").Parse(planAWSOverlay)
	FailIfError(err, "Couldn't parse template")

	By("Making infrastructure")
	etcdNode, etcErr := MakeETCDNode(nodeType)
	FailIfError(etcErr, "Error making etcd node")

	masterNode, masterErr := MakeMasterNode(nodeType)
	FailIfError(masterErr, "Error making master node")

	workerNode, workerErr := MakeWorkerNode(nodeType)
	FailIfError(workerErr, "Error making worker node")

	defer TerminateInstances(etcdNode.Instanceid, masterNode.Instanceid, workerNode.Instanceid)
	descErr := WaitForInstanceToStart(&etcdNode, &masterNode, &workerNode)
	FailIfError(descErr, "Error waiting for nodes")
	log.Printf("Created etcd nodes: %v (%v), master nodes %v (%v), workerNodes %v (%v)",
		etcdNode.Instanceid, etcdNode.Publicip,
		masterNode.Instanceid, masterNode.Publicip,
		workerNode.Instanceid, workerNode.Publicip)

	By("Building a plan to set up an overlay network cluster on this hardware")
	nodes := PlanUbuntuAWS{
		Etcd:                []AWSNodeDeets{etcdNode},
		Master:              []AWSNodeDeets{masterNode},
		Worker:              []AWSNodeDeets{workerNode},
		MasterNodeFQDN:      masterNode.Hostname,
		MasterNodeShortName: masterNode.Hostname,
		User:                user,
	}
	f, fileErr := os.Create("kismatic-testing.yaml")
	FailIfError(fileErr, "Error waiting for nodes")
	defer f.Close()
	w := bufio.NewWriter(f)
	execErr := template.Execute(w, &nodes)
	FailIfError(execErr, "Error filling in plan template")
	w.Flush()

	By("Validing our plan")
	ver := exec.Command("./kismatic", "install", "validate", "-f", f.Name())
	verbytes, verErr := ver.CombinedOutput()
	verText := string(verbytes)

	FailIfError(verErr, "Error validating plan", verText)

	By("Punch it Chewie!")
	app := exec.Command("./kismatic", "install", "apply", "-f", f.Name())
	appbytes, appErr := app.CombinedOutput()
	appText := string(appbytes)

	FailIfError(appErr, "Error applying plan", appText)
}

func FailIfError(err error, message ...string) {
	if err != nil {
		log.Printf(message[0]+": %v\n%v", err, message[1:])
		Fail(message[0])
	}
}

func CopyKismaticToTemp() string {
	tmpDir := os.TempDir()
	randomness, randomErr := GenerateGUIDString()
	if randomErr != nil {
		log.Fatal("Error making a GUID: ", randomErr)
	}
	kisPath := tmpDir + "/kisint/" + randomness
	err := os.MkdirAll(kisPath, 0777)
	if err != nil {
		log.Fatal("Error making temp dir: ", err)
	}

	return kisPath
}

func GenerateGUIDString() (string, error) {
	randomness, randomErr := guidMaker.NextId()

	if randomErr != nil {
		return "", randomErr
	}

	return strconv.FormatInt(randomness, 16), nil
}

func AssertKismaticDirectory(kisPath string) {
	if FileExists(kisPath + "/kismatic") {
		log.Fatal("Installer unpacked but kismatic wasn't there")
	}
}

func FileExists(path string) bool {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}

func MakeETCDNode(nodeType string) (AWSNodeDeets, error) {
	return MakeAWSNode(nodeType, ec2.InstanceTypeT2Micro)
}

func MakeMasterNode(nodeType string) (AWSNodeDeets, error) {
	return MakeAWSNode(nodeType, ec2.InstanceTypeT2Micro)
}

func MakeWorkerNode(nodeType string) (AWSNodeDeets, error) {
	return MakeAWSNode(nodeType, ec2.InstanceTypeT2Medium)
}

func MakeAWSNode(ami string, instanceType string) (AWSNodeDeets, error) {
	svc := ec2.New(session.New(&aws.Config{Region: aws.String(TARGET_REGION)}))
	runResult, err := svc.RunInstances(&ec2.RunInstancesInput{
		ImageId:          aws.String(ami),
		InstanceType:     aws.String(instanceType),
		MinCount:         aws.Int64(1),
		MaxCount:         aws.Int64(1),
		SubnetId:         aws.String(SUBNETID),
		KeyName:          aws.String(KEYNAME),
		SecurityGroupIds: []*string{aws.String(SECURITYGROUPID)},
	})

	if err != nil {
		return AWSNodeDeets{}, err
	}

	re := regexp.MustCompile("[^.]*")
	hostname := re.FindString(*runResult.Instances[0].PrivateDnsName)

	deets := AWSNodeDeets{
		Instanceid: *runResult.Instances[0].InstanceId,
		Privateip:  *runResult.Instances[0].PrivateIpAddress,
		Hostname:   hostname,
	}

	_, errtag := svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{aws.String(deets.Instanceid)},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("ApprendaTeam"),
				Value: aws.String("Kismatic"),
			},
		},
	})
	if errtag != nil {
		return deets, errtag
	}

	return deets, nil
}

func TerminateInstances(instanceids ...string) {
	awsinstanceids := make([]*string, len(instanceids))
	for i, id := range instanceids {
		awsinstanceids[i] = aws.String(id)
	}
	sess, err := session.NewSession()

	if err != nil {
		log.Printf("failed to create session: %v", err)
		return
	}

	svc := ec2.New(sess, &aws.Config{Region: aws.String(TARGET_REGION)})

	params := &ec2.TerminateInstancesInput{
		InstanceIds: awsinstanceids,
	}
	resp, err := svc.TerminateInstances(params)

	if err != nil {
		log.Printf("Could not terminate: %v", resp)
		return
	}
}

func WaitForInstanceToStart(nodes ...*AWSNodeDeets) error {
	sess, err := session.NewSession()

	if err != nil {
		fmt.Println("failed to create session,", err)
		return err
	}

	fmt.Print("Waiting for nodes to come up")
	defer fmt.Println()

	svc := ec2.New(sess, &aws.Config{Region: aws.String(TARGET_REGION)})
	for _, deets := range nodes {
		deets.Publicip = ""

		for deets.Publicip == "" {
			fmt.Print(".")
			descResult, descErr := svc.DescribeInstances(&ec2.DescribeInstancesInput{
				InstanceIds: []*string{aws.String(deets.Instanceid)},
			})
			if descErr != nil {
				return descErr
			}

			if *descResult.Reservations[0].Instances[0].State.Name == ec2.InstanceStateNameRunning &&
				descResult.Reservations[0].Instances[0].PublicIpAddress != nil {
				deets.Publicip = *descResult.Reservations[0].Instances[0].PublicIpAddress
				BlockUntilSSHOpen(deets)
			} else {
				time.Sleep(1 * time.Second)
			}
		}
	}
	return nil
}

func BlockUntilSSHOpen(node *AWSNodeDeets) {
	conn, err := net.Dial("tcp", node.Publicip+":22")
	fmt.Print("?")
	if err != nil {
		time.Sleep(5 & time.Second)
		BlockUntilSSHOpen(node)
	} else {
		conn.Close()
		return
	}
}
<|endoftext|>