<commit_before>package proxy

import (
	"encoding/json"
	"strings"

	"github.com/asaskevich/govalidator"
	"github.com/hellofresh/janus/pkg/router"
)

// Definition defines proxy rules for a route
type Definition struct {
	PreserveHost        bool       `bson:"preserve_host" json:"preserve_host" mapstructure:"preserve_host"`
	ListenPath          string     `bson:"listen_path" json:"listen_path" mapstructure:"listen_path" valid:"required~proxy.listen_path is required,urlpath"`
	Upstreams           *Upstreams `bson:"upstreams" json:"upstreams" mapstructure:"upstreams"`
	InsecureSkipVerify  bool       `bson:"insecure_skip_verify" json:"insecure_skip_verify" mapstructure:"insecure_skip_verify"`
	StripPath           bool       `bson:"strip_path" json:"strip_path" mapstructure:"strip_path"`
	AppendPath          bool       `bson:"append_path" json:"append_path" mapstructure:"append_path"`
	EnableLoadBalancing bool       `bson:"enable_load_balancing" json:"enable_load_balancing" mapstructure:"enable_load_balancing"`
	Methods             []string   `bson:"methods" json:"methods"`
	Hosts               []string   `bson:"hosts" json:"hosts"`
}

// Upstreams represents a collection of targets where the requests will go to
type Upstreams struct {
	Balancing string    `bson:"balancing" json:"balancing"`
	Targets   []*Target `bson:"targets" json:"targets"`
}

// Target is an ip address/hostname with a port that identifies an instance of a backend service
type Target struct {
	Target string `bson:"target" json:"target" valid:"url,required"`
	Weight int    `bson:"weight" json:"weight"`
}

// NewDefinition creates a new Proxy Definition with default values
func NewDefinition() *Definition {
	return &Definition{
		Methods: make([]string, 0),
		Hosts:   make([]string, 0),
		Upstreams: &Upstreams{
			Targets: make([]*Target, 0),
		},
	}
}

// Validate validates proxy data
func (d *Definition) Validate() (bool, error) {
	return govalidator.ValidateStruct(d)
}

// IsBalancerDefined checks if load balancer is defined
func (d *Definition) IsBalancerDefined() bool {
	return d.Upstreams != nil && d.Upstreams.Targets != nil && len(d.Upstreams.Targets) > 0
}

// Route is the container for a proxy and it's handlers
type Route struct {
	Proxy    *Definition
	Inbound  InChain
	Outbound OutChain
}

type routeJSONProxy struct {
	Proxy *Definition `json:"proxy"`
}

// NewRoute creates an instance of Route
func NewRoute(proxy *Definition) *Route {
	return &Route{Proxy: proxy}
}

// NewRouteWithInOut creates an instance of Route with inbound and outbound handlers
func NewRouteWithInOut(proxy *Definition, inbound InChain, outbound OutChain) *Route {
	return &Route{proxy, inbound, outbound}
}

// AddInbound adds inbound middlewares
func (r *Route) AddInbound(in ...router.Constructor) {
	for _, i := range in {
		r.Inbound = append(r.Inbound, i)
	}
}

// AddOutbound adds outbound middlewares
func (r *Route) AddOutbound(out ...OutLink) {
	for _, o := range out {
		r.Outbound = append(r.Outbound, o)
	}
}

// JSONMarshal encodes route struct to JSON
func (r *Route) JSONMarshal() ([]byte, error) {
	return json.Marshal(routeJSONProxy{r.Proxy})
}

// JSONUnmarshalRoute decodes route struct from JSON
func JSONUnmarshalRoute(rawRoute []byte) (*Route, error) {
	var proxyRoute routeJSONProxy
	if err := json.Unmarshal(rawRoute, &proxyRoute); err != nil {
		return nil, err
	}
	return NewRoute(proxyRoute.Proxy), nil
}

func init() {
	// initializes custom validators
	govalidator.CustomTypeTagMap.Set("urlpath", func(i interface{}, o interface{}) bool {
		s, ok := i.(string)
		if !ok {
			return false
		}

		return strings.Index(s, "/") == 0
	})
}
<commit_msg>Fix default method<commit_after>package proxy

import (
	"encoding/json"
	"strings"

	"github.com/asaskevich/govalidator"
	"github.com/hellofresh/janus/pkg/router"
)

// Definition defines proxy rules for a route
type Definition struct {
	PreserveHost        bool       `bson:"preserve_host" json:"preserve_host" mapstructure:"preserve_host"`
	ListenPath          string     `bson:"listen_path" json:"listen_path" mapstructure:"listen_path" valid:"required~proxy.listen_path is required,urlpath"`
	Upstreams           *Upstreams `bson:"upstreams" json:"upstreams" mapstructure:"upstreams"`
	InsecureSkipVerify  bool       `bson:"insecure_skip_verify" json:"insecure_skip_verify" mapstructure:"insecure_skip_verify"`
	StripPath           bool       `bson:"strip_path" json:"strip_path" mapstructure:"strip_path"`
	AppendPath          bool       `bson:"append_path" json:"append_path" mapstructure:"append_path"`
	EnableLoadBalancing bool       `bson:"enable_load_balancing" json:"enable_load_balancing" mapstructure:"enable_load_balancing"`
	Methods             []string   `bson:"methods" json:"methods"`
	Hosts               []string   `bson:"hosts" json:"hosts"`
}

// Upstreams represents a collection of targets where the requests will go to
type Upstreams struct {
	Balancing string    `bson:"balancing" json:"balancing"`
	Targets   []*Target `bson:"targets" json:"targets"`
}

// Target is an ip address/hostname with a port that identifies an instance of a backend service
type Target struct {
	Target string `bson:"target" json:"target" valid:"url,required"`
	Weight int    `bson:"weight" json:"weight"`
}

// NewDefinition creates a new Proxy Definition with default values
func NewDefinition() *Definition {
	return &Definition{
		Methods: []string{"GET"},
		Hosts:   make([]string, 0),
		Upstreams: &Upstreams{
			Targets: make([]*Target, 0),
		},
	}
}

// Validate validates proxy data
func (d *Definition) Validate() (bool, error) {
	return govalidator.ValidateStruct(d)
}

// IsBalancerDefined checks if load balancer is defined
func (d *Definition) IsBalancerDefined() bool {
	return d.Upstreams != nil && d.Upstreams.Targets != nil && len(d.Upstreams.Targets) > 0
}

// Route is the container for a proxy and it's handlers
type Route struct {
	Proxy    *Definition
	Inbound  InChain
	Outbound OutChain
}

type routeJSONProxy struct {
	Proxy *Definition `json:"proxy"`
}

// NewRoute creates an instance of Route
func NewRoute(proxy *Definition) *Route {
	return &Route{Proxy: proxy}
}

// NewRouteWithInOut creates an instance of Route with inbound and outbound handlers
func NewRouteWithInOut(proxy *Definition, inbound InChain, outbound OutChain) *Route {
	return &Route{proxy, inbound, outbound}
}

// AddInbound adds inbound middlewares
func (r *Route) AddInbound(in ...router.Constructor) {
	for _, i := range in {
		r.Inbound = append(r.Inbound, i)
	}
}

// AddOutbound adds outbound middlewares
func (r *Route) AddOutbound(out ...OutLink) {
	for _, o := range out {
		r.Outbound = append(r.Outbound, o)
	}
}

// JSONMarshal encodes route struct to JSON
func (r *Route) JSONMarshal() ([]byte, error) {
	return json.Marshal(routeJSONProxy{r.Proxy})
}

// JSONUnmarshalRoute decodes route struct from JSON
func JSONUnmarshalRoute(rawRoute []byte) (*Route, error) {
	var proxyRoute routeJSONProxy
	if err := json.Unmarshal(rawRoute, &proxyRoute); err != nil {
		return nil, err
	}
	return NewRoute(proxyRoute.Proxy), nil
}

func init() {
	// initializes custom validators
	govalidator.CustomTypeTagMap.Set("urlpath", func(i interface{}, o interface{}) bool {
		s, ok := i.(string)
		if !ok {
			return false
		}

		return strings.Index(s, "/") == 0
	})
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"strings"
	"time"

	"github.com/abh/geodns/Godeps/_workspace/src/github.com/miekg/dns"
	. "github.com/abh/geodns/Godeps/_workspace/src/gopkg.in/check.v1"
)

const (
	PORT = ":8853"
)

type ServeSuite struct {
}

var _ = Suite(&ServeSuite{})

func (s *ServeSuite) SetUpSuite(c *C) {

	// setup and register metrics
	metrics := NewMetrics()
	go metrics.Updater()

	Zones := make(Zones)
	setupPgeodnsZone(Zones)
	setupRootZone()
	zonesReadDir("dns", Zones)

	go listenAndServe(PORT)

	time.Sleep(200 * time.Millisecond)
}

func (s *ServeSuite) TestServing(c *C) {

	r := exchange(c, "_status.pgeodns.", dns.TypeTXT)
	txt := r.Answer[0].(*dns.TXT).Txt[0]
	if !strings.HasPrefix(txt, "{") {
		c.Log("Unexpected result for _status.pgeodns", txt)
		c.Fail()
	}

	// Allow _country and _status queries as long as the first label is that
	r = exchange(c, "_country.foo.pgeodns.", dns.TypeTXT)
	txt = r.Answer[0].(*dns.TXT).Txt[0]
	// Got appropriate response for _country txt query
	if !strings.HasPrefix(txt, "127.0.0.1:") {
		c.Log("Unexpected result for _country.foo.pgeodns", txt)
		c.Fail()
	}

	// Make sure A requests for _status doesn't NXDOMAIN
	r = exchange(c, "_status.pgeodns.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	// Got one SOA record
	c.Check(r.Ns, HasLen, 1)
	// NOERROR for A request
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	r = exchange(c, "bar.test.example.com.", dns.TypeA)
	ip := r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.2")
	c.Check(int(r.Answer[0].Header().Ttl), Equals, 601)

	r = exchange(c, "test.example.com.", dns.TypeSOA)
	soa := r.Answer[0].(*dns.SOA)
	serial := soa.Serial
	c.Check(int(serial), Equals, 3)

	// no AAAA records for 'bar', so check we get a soa record back
	r = exchange(c, "test.example.com.", dns.TypeAAAA)
	soa2 := r.Ns[0].(*dns.SOA)
	c.Check(soa, DeepEquals, soa2)

	// CNAMEs
	r = exchange(c, "www.test.example.com.", dns.TypeA)
	c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo.bitnames.com.")
	c.Check(int(r.Answer[0].Header().Ttl), Equals, 1800)

	//SPF
	r = exchange(c, "test.example.com.", dns.TypeSPF)
	c.Check(r.Answer[0].(*dns.SPF).Txt[0], Equals, "v=spf1 ~all")

	//SRV
	r = exchange(c, "_sip._tcp.test.example.com.", dns.TypeSRV)
	c.Check(r.Answer[0].(*dns.SRV).Target, Equals, "sipserver.example.com.")
	c.Check(r.Answer[0].(*dns.SRV).Port, Equals, uint16(5060))
	c.Check(r.Answer[0].(*dns.SRV).Priority, Equals, uint16(10))
	c.Check(r.Answer[0].(*dns.SRV).Weight, Equals, uint16(100))

	// MX
	r = exchange(c, "test.example.com.", dns.TypeMX)
	c.Check(r.Answer[0].(*dns.MX).Mx, Equals, "mx.example.net.")
	c.Check(r.Answer[1].(*dns.MX).Mx, Equals, "mx2.example.net.")
	c.Check(r.Answer[1].(*dns.MX).Preference, Equals, uint16(20))

	// Verify the first A record was created
	r = exchange(c, "a.b.c.test.example.com.", dns.TypeA)
	ip = r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.7")

	// Verify sub-labels are created
	r = exchange(c, "b.c.test.example.com.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	r = exchange(c, "c.test.example.com.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	// Verify the first A record was created
	r = exchange(c, "three.two.one.test.example.com.", dns.TypeA)
	ip = r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.5")

	// Verify single sub-labels is created and no record is returned
	r = exchange(c, "two.one.test.example.com.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	// Verify the A record wasn't over written
	r = exchange(c, "one.test.example.com.", dns.TypeA)
	ip = r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.6")
}

func (s *ServeSuite) TestServingMixedCase(c *C) {

	r := exchange(c, "_sTaTUs.pGEOdns.", dns.TypeTXT)
	c.Assert(r.Rcode, Equals, dns.RcodeSuccess)
	txt := r.Answer[0].(*dns.TXT).Txt[0]
	if !strings.HasPrefix(txt, "{") {
		c.Log("Unexpected result for _status.pgeodns", txt)
		c.Fail()
	}

	n := "baR.test.eXAmPLe.cOM."
	r = exchange(c, n, dns.TypeA)
	ip := r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.2")
	c.Check(r.Answer[0].Header().Name, Equals, n)

}

func (s *ServeSuite) TestCname(c *C) {
	// Cname, two possible results

	results := make(map[string]int)

	for i := 0; i < 10; i++ {
		r := exchange(c, "www.se.test.example.com.", dns.TypeA)
		target := r.Answer[0].(*dns.CNAME).Target
		results[target]++
	}

	// Two possible results from this cname
	c.Check(results, HasLen, 2)
}

func (s *ServeSuite) TestUnknownDomain(c *C) {
	r := exchange(c, "no.such.domain.", dns.TypeAAAA)
	c.Assert(r.Rcode, Equals, dns.RcodeRefused)
}

func (s *ServeSuite) TestServingAliases(c *C) {
	// Alias, no geo matches
	r := exchange(c, "bar-alias.test.example.com.", dns.TypeA)
	ip := r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.2")

	// Alias to a cname record
	r = exchange(c, "www-alias.test.example.com.", dns.TypeA)
	c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo.bitnames.com.")

	// Alias returning a cname, with geo overrides
	r = exchangeSubnet(c, "www-alias.test.example.com.", dns.TypeA, "194.239.134.1")
	c.Check(r.Answer, HasLen, 1)
	if len(r.Answer) > 0 {
		c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo-europe.bitnames.com.")
	}

	// Alias to Ns records
	r = exchange(c, "sub-alias.test.example.org.", dns.TypeNS)
	c.Check(r.Answer[0].(*dns.NS).Ns, Equals, "ns1.example.com.")

}

func (s *ServeSuite) TestServingEDNS(c *C) {
	// MX test
	r := exchangeSubnet(c, "test.example.com.", dns.TypeMX, "194.239.134.1")
	c.Check(r.Answer, HasLen, 1)
	if len(r.Answer) > 0 {
		c.Check(r.Answer[0].(*dns.MX).Mx, Equals, "mx-eu.example.net.")
	}

	c.Log("Testing www.test.example.com from .dk, should match www.europe (a cname)")

	r = exchangeSubnet(c, "www.test.example.com.", dns.TypeA, "194.239.134.0")
	// www.test from .dk IP address gets at least one answer
	c.Check(r.Answer, HasLen, 1)
	if len(r.Answer) > 0 {
		// EDNS-SUBNET test (request A, respond CNAME)
		c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo-europe.bitnames.com.")
	}

}

func (s *ServeSuite) BenchmarkServing(c *C) {
	for i := 0; i < c.N; i++ {
		exchange(c, "_country.foo.pgeodns.", dns.TypeTXT)
	}
}

func exchangeSubnet(c *C, name string, dnstype uint16, ip string) *dns.Msg {
	msg := new(dns.Msg)

	msg.SetQuestion(name, dnstype)

	o := new(dns.OPT)
	o.Hdr.Name = "."
	o.Hdr.Rrtype = dns.TypeOPT
	e := new(dns.EDNS0_SUBNET)
	e.Code = dns.EDNS0SUBNET
	e.SourceScope = 0
	e.Address = net.ParseIP(ip)
	e.Family = 1 // IP4
	e.SourceNetmask = net.IPv4len * 8
	o.Option = append(o.Option, e)
	msg.Extra = append(msg.Extra, o)

	c.Log("msg", msg)

	return dorequest(c, msg)
}

func exchange(c *C, name string, dnstype uint16) *dns.Msg {
	msg := new(dns.Msg)

	msg.SetQuestion(name, dnstype)
	return dorequest(c, msg)
}

func dorequest(c *C, msg *dns.Msg) *dns.Msg {
	cli := new(dns.Client)
	r, _, err := cli.Exchange(msg, "127.0.0.1"+PORT)
	if err != nil {
		c.Log("err", err)
		c.Fail()
	}
	return r
}
<commit_msg>Add test to find data races<commit_after>package main

import (
	"net"
	"strings"
	"sync"
	"time"

	"github.com/abh/geodns/Godeps/_workspace/src/github.com/miekg/dns"
	. "github.com/abh/geodns/Godeps/_workspace/src/gopkg.in/check.v1"
)

const (
	PORT = ":8853"
)

type ServeSuite struct {
}

var _ = Suite(&ServeSuite{})

func (s *ServeSuite) SetUpSuite(c *C) {

	// setup and register metrics
	metrics := NewMetrics()
	go metrics.Updater()

	Zones := make(Zones)
	setupPgeodnsZone(Zones)
	setupRootZone()
	zonesReadDir("dns", Zones)

	go listenAndServe(PORT)

	time.Sleep(200 * time.Millisecond)
}

func (s *ServeSuite) TestServing(c *C) {

	r := exchange(c, "_status.pgeodns.", dns.TypeTXT)
	txt := r.Answer[0].(*dns.TXT).Txt[0]
	if !strings.HasPrefix(txt, "{") {
		c.Log("Unexpected result for _status.pgeodns", txt)
		c.Fail()
	}

	// Allow _country and _status queries as long as the first label is that
	r = exchange(c, "_country.foo.pgeodns.", dns.TypeTXT)
	txt = r.Answer[0].(*dns.TXT).Txt[0]
	// Got appropriate response for _country txt query
	if !strings.HasPrefix(txt, "127.0.0.1:") {
		c.Log("Unexpected result for _country.foo.pgeodns", txt)
		c.Fail()
	}

	// Make sure A requests for _status doesn't NXDOMAIN
	r = exchange(c, "_status.pgeodns.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	// Got one SOA record
	c.Check(r.Ns, HasLen, 1)
	// NOERROR for A request
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	r = exchange(c, "bar.test.example.com.", dns.TypeA)
	ip := r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.2")
	c.Check(int(r.Answer[0].Header().Ttl), Equals, 601)

	r = exchange(c, "test.example.com.", dns.TypeSOA)
	soa := r.Answer[0].(*dns.SOA)
	serial := soa.Serial
	c.Check(int(serial), Equals, 3)

	// no AAAA records for 'bar', so check we get a soa record back
	r = exchange(c, "test.example.com.", dns.TypeAAAA)
	soa2 := r.Ns[0].(*dns.SOA)
	c.Check(soa, DeepEquals, soa2)

	// CNAMEs
	r = exchange(c, "www.test.example.com.", dns.TypeA)
	c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo.bitnames.com.")
	c.Check(int(r.Answer[0].Header().Ttl), Equals, 1800)

	//SPF
	r = exchange(c, "test.example.com.", dns.TypeSPF)
	c.Check(r.Answer[0].(*dns.SPF).Txt[0], Equals, "v=spf1 ~all")

	//SRV
	r = exchange(c, "_sip._tcp.test.example.com.", dns.TypeSRV)
	c.Check(r.Answer[0].(*dns.SRV).Target, Equals, "sipserver.example.com.")
	c.Check(r.Answer[0].(*dns.SRV).Port, Equals, uint16(5060))
	c.Check(r.Answer[0].(*dns.SRV).Priority, Equals, uint16(10))
	c.Check(r.Answer[0].(*dns.SRV).Weight, Equals, uint16(100))

	// MX
	r = exchange(c, "test.example.com.", dns.TypeMX)
	c.Check(r.Answer[0].(*dns.MX).Mx, Equals, "mx.example.net.")
	c.Check(r.Answer[1].(*dns.MX).Mx, Equals, "mx2.example.net.")
	c.Check(r.Answer[1].(*dns.MX).Preference, Equals, uint16(20))

	// Verify the first A record was created
	r = exchange(c, "a.b.c.test.example.com.", dns.TypeA)
	ip = r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.7")

	// Verify sub-labels are created
	r = exchange(c, "b.c.test.example.com.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	r = exchange(c, "c.test.example.com.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	// Verify the first A record was created
	r = exchange(c, "three.two.one.test.example.com.", dns.TypeA)
	ip = r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.5")

	// Verify single sub-labels is created and no record is returned
	r = exchange(c, "two.one.test.example.com.", dns.TypeA)
	c.Check(r.Answer, HasLen, 0)
	c.Check(r.Rcode, Equals, dns.RcodeSuccess)

	// Verify the A record wasn't over written
	r = exchange(c, "one.test.example.com.", dns.TypeA)
	ip = r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.6")
}

func (s *ServeSuite) TestServingMixedCase(c *C) {

	r := exchange(c, "_sTaTUs.pGEOdns.", dns.TypeTXT)
	c.Assert(r.Rcode, Equals, dns.RcodeSuccess)
	txt := r.Answer[0].(*dns.TXT).Txt[0]
	if !strings.HasPrefix(txt, "{") {
		c.Log("Unexpected result for _status.pgeodns", txt)
		c.Fail()
	}

	n := "baR.test.eXAmPLe.cOM."
	r = exchange(c, n, dns.TypeA)
	ip := r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.2")
	c.Check(r.Answer[0].Header().Name, Equals, n)

}

func (s *ServeSuite) TestCname(c *C) {
	// Cname, two possible results

	results := make(map[string]int)

	for i := 0; i < 10; i++ {
		r := exchange(c, "www.se.test.example.com.", dns.TypeA)
		target := r.Answer[0].(*dns.CNAME).Target
		results[target]++
	}

	// Two possible results from this cname
	c.Check(results, HasLen, 2)
}

func (s *ServeSuite) TestUnknownDomain(c *C) {
	r := exchange(c, "no.such.domain.", dns.TypeAAAA)
	c.Assert(r.Rcode, Equals, dns.RcodeRefused)
}

func (s *ServeSuite) TestServingAliases(c *C) {
	// Alias, no geo matches
	r := exchange(c, "bar-alias.test.example.com.", dns.TypeA)
	ip := r.Answer[0].(*dns.A).A
	c.Check(ip.String(), Equals, "192.168.1.2")

	// Alias to a cname record
	r = exchange(c, "www-alias.test.example.com.", dns.TypeA)
	c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo.bitnames.com.")

	// Alias returning a cname, with geo overrides
	r = exchangeSubnet(c, "www-alias.test.example.com.", dns.TypeA, "194.239.134.1")
	c.Check(r.Answer, HasLen, 1)
	if len(r.Answer) > 0 {
		c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo-europe.bitnames.com.")
	}

	// Alias to Ns records
	r = exchange(c, "sub-alias.test.example.org.", dns.TypeNS)
	c.Check(r.Answer[0].(*dns.NS).Ns, Equals, "ns1.example.com.")

}

func (s *ServeSuite) TestServingEDNS(c *C) {
	// MX test
	r := exchangeSubnet(c, "test.example.com.", dns.TypeMX, "194.239.134.1")
	c.Check(r.Answer, HasLen, 1)
	if len(r.Answer) > 0 {
		c.Check(r.Answer[0].(*dns.MX).Mx, Equals, "mx-eu.example.net.")
	}

	c.Log("Testing www.test.example.com from .dk, should match www.europe (a cname)")

	r = exchangeSubnet(c, "www.test.example.com.", dns.TypeA, "194.239.134.0")
	// www.test from .dk IP address gets at least one answer
	c.Check(r.Answer, HasLen, 1)
	if len(r.Answer) > 0 {
		// EDNS-SUBNET test (request A, respond CNAME)
		c.Check(r.Answer[0].(*dns.CNAME).Target, Equals, "geo-europe.bitnames.com.")
	}

}

func (s *ServeSuite) TestServeRace(c *C) {
	wg := sync.WaitGroup{}
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			s.TestServing(c)
			wg.Done()
		}()
	}
	wg.Wait()
}

func (s *ServeSuite) BenchmarkServing(c *C) {
	for i := 0; i < c.N; i++ {
		exchange(c, "_country.foo.pgeodns.", dns.TypeTXT)
	}
}

func exchangeSubnet(c *C, name string, dnstype uint16, ip string) *dns.Msg {
	msg := new(dns.Msg)

	msg.SetQuestion(name, dnstype)

	o := new(dns.OPT)
	o.Hdr.Name = "."
	o.Hdr.Rrtype = dns.TypeOPT
	e := new(dns.EDNS0_SUBNET)
	e.Code = dns.EDNS0SUBNET
	e.SourceScope = 0
	e.Address = net.ParseIP(ip)
	e.Family = 1 // IP4
	e.SourceNetmask = net.IPv4len * 8
	o.Option = append(o.Option, e)
	msg.Extra = append(msg.Extra, o)

	c.Log("msg", msg)

	return dorequest(c, msg)
}

func exchange(c *C, name string, dnstype uint16) *dns.Msg {
	msg := new(dns.Msg)

	msg.SetQuestion(name, dnstype)
	return dorequest(c, msg)
}

func dorequest(c *C, msg *dns.Msg) *dns.Msg {
	cli := new(dns.Client)
	r, _, err := cli.Exchange(msg, "127.0.0.1"+PORT)
	if err != nil {
		c.Log("err", err)
		c.Fail()
	}
	return r
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package estimator

import (
	"fmt"
	"sort"
	"time"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/autoscaler/cluster-autoscaler/simulator"
	"k8s.io/klog"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
)

// podInfo contains Pod and score that corresponds to how important it is to handle the pod first.
type podInfo struct {
	score float64
	pod   *apiv1.Pod
}

// BinpackingNodeEstimator estimates the number of needed nodes to handle the given amount of pods.
type BinpackingNodeEstimator struct {
	predicateChecker simulator.PredicateChecker
	clusterSnapshot  simulator.ClusterSnapshot
}

// NewBinpackingNodeEstimator builds a new BinpackingNodeEstimator.
func NewBinpackingNodeEstimator(
	predicateChecker simulator.PredicateChecker,
	clusterSnapshot simulator.ClusterSnapshot) *BinpackingNodeEstimator {
	return &BinpackingNodeEstimator{
		predicateChecker: predicateChecker,
		clusterSnapshot:  clusterSnapshot,
	}
}

// Estimate implements First Fit Decreasing bin-packing approximation algorithm.
// See https://en.wikipedia.org/wiki/Bin_packing_problem for more details.
// While it is a multi-dimensional bin packing (cpu, mem, ports) in most cases the main dimension
// will be cpu thus the estimated overprovisioning of 11/9 * optimal + 6/9 should be
// still be maintained.
// It is assumed that all pods from the given list can fit to nodeTemplate.
// Returns the number of nodes needed to accommodate all pods from the list.
func (estimator *BinpackingNodeEstimator) Estimate(
	pods []*apiv1.Pod,
	nodeTemplate *schedulernodeinfo.NodeInfo,
	upcomingNodes []*schedulernodeinfo.NodeInfo) int {
	podInfos := calculatePodScore(pods, nodeTemplate)
	sort.Slice(podInfos, func(i, j int) bool { return podInfos[i].score > podInfos[j].score })

	newNodeNames := make([]string, 0)

	if err := estimator.clusterSnapshot.Fork(); err != nil {
		klog.Errorf("Error while calling ClusterSnapshot.Fork; %v", err)
		return 0
	}
	defer func() {
		if err := estimator.clusterSnapshot.Revert(); err != nil {
			klog.Errorf("Error while calling ClusterSnapshot.Revert; %v", err)
		}
	}()

	newNodeNameTimestamp := time.Now()
	newNodeNameIndex := 0

	for _, podInfo := range podInfos {
		found := false
		for _, nodeName := range newNodeNames {
			if err := estimator.predicateChecker.CheckPredicates(estimator.clusterSnapshot, podInfo.pod, simulator.FakeNodeInfoForNodeName(nodeName)); err == nil {
				found = true
				if err := estimator.clusterSnapshot.AddPod(podInfo.pod, nodeName); err != nil {
					klog.Errorf("Error adding pod %v.%v to node %v in ClusterSnapshot; %v", podInfo.pod.Namespace, podInfo.pod.Name, nodeName, err)
					return 0
				}
				break
			}
		}
		if !found {
			// Add new node
			newNodeName, err := estimator.addNewNodeToSnapshot(nodeTemplate, newNodeNameTimestamp, newNodeNameIndex)
			if err != nil {
				klog.Errorf("Error while adding new node for template to ClusterSnapshot; %v", err)
				return 0
			}
			newNodeNameIndex++
			// And schedule pod to it
			if err := estimator.clusterSnapshot.AddPod(podInfo.pod, newNodeName); err != nil {
				klog.Errorf("Error adding pod %v.%v to node %v in ClusterSnapshot; %v", podInfo.pod.Namespace, podInfo.pod.Name, newNodeName, err)
				return 0
			}
			newNodeNames = append(newNodeNames, newNodeName)
		}
	}
	return len(newNodeNames)
}

func (estimator *BinpackingNodeEstimator) addNewNodeToSnapshot(
	template *schedulernodeinfo.NodeInfo,
	nameTimestamp time.Time,
	nameIndex int) (string, error) {

	newNode := template.Node().DeepCopy()
	newNode.Name = fmt.Sprintf("%s-%d-%d", newNode.Name, nameTimestamp.Unix(), nameIndex)
	if err := estimator.clusterSnapshot.AddNodeWithPods(newNode, template.Pods()); err != nil {
		return "", err
	}
	return newNode.Name, nil
}

// Calculates score for all pods and returns podInfo structure.
// Score is defined as cpu_sum/node_capacity + mem_sum/node_capacity.
// Pods that have bigger requirements should be processed first, thus have higher scores.
func calculatePodScore(pods []*apiv1.Pod, nodeTemplate *schedulernodeinfo.NodeInfo) []*podInfo {
	podInfos := make([]*podInfo, 0, len(pods))

	for _, pod := range pods {
		cpuSum := resource.Quantity{}
		memorySum := resource.Quantity{}

		for _, container := range pod.Spec.Containers {
			if request, ok := container.Resources.Requests[apiv1.ResourceCPU]; ok {
				cpuSum.Add(request)
			}
			if request, ok := container.Resources.Requests[apiv1.ResourceMemory]; ok {
				memorySum.Add(request)
			}
		}
		score := float64(0)
		if cpuAllocatable, ok := nodeTemplate.Node().Status.Allocatable[apiv1.ResourceCPU]; ok && cpuAllocatable.MilliValue() > 0 {
			score += float64(cpuSum.MilliValue()) / float64(cpuAllocatable.MilliValue())
		}
		if memAllocatable, ok := nodeTemplate.Node().Status.Allocatable[apiv1.ResourceMemory]; ok && memAllocatable.Value() > 0 {
			score += float64(memorySum.Value()) / float64(memAllocatable.Value())
		}

		podInfos = append(podInfos, &podInfo{
			score: score,
			pod:   pod,
		})
	}
	return podInfos
}
<commit_msg>Set kubernetes.io/hostname label on simulated node in BinpackingNodeEstimator<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package estimator

import (
	"fmt"
	"sort"
	"time"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/autoscaler/cluster-autoscaler/simulator"
	"k8s.io/klog"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
)

// podInfo contains Pod and score that corresponds to how important it is to handle the pod first.
type podInfo struct {
	score float64
	pod   *apiv1.Pod
}

// BinpackingNodeEstimator estimates the number of needed nodes to handle the given amount of pods.
type BinpackingNodeEstimator struct {
	predicateChecker simulator.PredicateChecker
	clusterSnapshot  simulator.ClusterSnapshot
}

// NewBinpackingNodeEstimator builds a new BinpackingNodeEstimator.
func NewBinpackingNodeEstimator(
	predicateChecker simulator.PredicateChecker,
	clusterSnapshot simulator.ClusterSnapshot) *BinpackingNodeEstimator {
	return &BinpackingNodeEstimator{
		predicateChecker: predicateChecker,
		clusterSnapshot:  clusterSnapshot,
	}
}

// Estimate implements First Fit Decreasing bin-packing approximation algorithm.
// See https://en.wikipedia.org/wiki/Bin_packing_problem for more details.
// While it is a multi-dimensional bin packing (cpu, mem, ports) in most cases the main dimension
// will be cpu thus the estimated overprovisioning of 11/9 * optimal + 6/9 should be
// still be maintained.
// It is assumed that all pods from the given list can fit to nodeTemplate.
// Returns the number of nodes needed to accommodate all pods from the list.
func (estimator *BinpackingNodeEstimator) Estimate(
	pods []*apiv1.Pod,
	nodeTemplate *schedulernodeinfo.NodeInfo,
	upcomingNodes []*schedulernodeinfo.NodeInfo) int {
	podInfos := calculatePodScore(pods, nodeTemplate)
	sort.Slice(podInfos, func(i, j int) bool { return podInfos[i].score > podInfos[j].score })

	newNodeNames := make([]string, 0)

	if err := estimator.clusterSnapshot.Fork(); err != nil {
		klog.Errorf("Error while calling ClusterSnapshot.Fork; %v", err)
		return 0
	}
	defer func() {
		if err := estimator.clusterSnapshot.Revert(); err != nil {
			klog.Errorf("Error while calling ClusterSnapshot.Revert; %v", err)
		}
	}()

	newNodeNameTimestamp := time.Now()
	newNodeNameIndex := 0

	for _, podInfo := range podInfos {
		found := false
		for _, nodeName := range newNodeNames {
			if err := estimator.predicateChecker.CheckPredicates(estimator.clusterSnapshot, podInfo.pod, simulator.FakeNodeInfoForNodeName(nodeName)); err == nil {
				found = true
				if err := estimator.clusterSnapshot.AddPod(podInfo.pod, nodeName); err != nil {
					klog.Errorf("Error adding pod %v.%v to node %v in ClusterSnapshot; %v", podInfo.pod.Namespace, podInfo.pod.Name, nodeName, err)
					return 0
				}
				break
			}
		}
		if !found {
			// Add new node
			newNodeName, err := estimator.addNewNodeToSnapshot(nodeTemplate, newNodeNameTimestamp, newNodeNameIndex)
			if err != nil {
				klog.Errorf("Error while adding new node for template to ClusterSnapshot; %v", err)
				return 0
			}
			newNodeNameIndex++
			// And schedule pod to it
			if err := estimator.clusterSnapshot.AddPod(podInfo.pod, newNodeName); err != nil {
				klog.Errorf("Error adding pod %v.%v to node %v in ClusterSnapshot; %v", podInfo.pod.Namespace, podInfo.pod.Name, newNodeName, err)
				return 0
			}
			newNodeNames = append(newNodeNames, newNodeName)
		}
	}
	return len(newNodeNames)
}

func (estimator *BinpackingNodeEstimator) addNewNodeToSnapshot(
	template *schedulernodeinfo.NodeInfo,
	nameTimestamp time.Time,
	nameIndex int) (string, error) {

	newNode := template.Node().DeepCopy()
	newNode.Name = fmt.Sprintf("%s-%d-%d", newNode.Name, nameTimestamp.Unix(), nameIndex)
	if newNode.Labels == nil {
		newNode.Labels = make(map[string]string)
	}
	newNode.Labels["kubernetes.io/hostname"] = newNode.Name
	if err := estimator.clusterSnapshot.AddNodeWithPods(newNode, template.Pods()); err != nil {
		return "", err
	}
	return newNode.Name, nil
}

// Calculates score for all pods and returns podInfo structure.
// Score is defined as cpu_sum/node_capacity + mem_sum/node_capacity.
// Pods that have bigger requirements should be processed first, thus have higher scores.
func calculatePodScore(pods []*apiv1.Pod, nodeTemplate *schedulernodeinfo.NodeInfo) []*podInfo {
	podInfos := make([]*podInfo, 0, len(pods))

	for _, pod := range pods {
		cpuSum := resource.Quantity{}
		memorySum := resource.Quantity{}

		for _, container := range pod.Spec.Containers {
			if request, ok := container.Resources.Requests[apiv1.ResourceCPU]; ok {
				cpuSum.Add(request)
			}
			if request, ok := container.Resources.Requests[apiv1.ResourceMemory]; ok {
				memorySum.Add(request)
			}
		}
		score := float64(0)
		if cpuAllocatable, ok := nodeTemplate.Node().Status.Allocatable[apiv1.ResourceCPU]; ok && cpuAllocatable.MilliValue() > 0 {
			score += float64(cpuSum.MilliValue()) / float64(cpuAllocatable.MilliValue())
		}
		if memAllocatable, ok := nodeTemplate.Node().Status.Allocatable[apiv1.ResourceMemory]; ok && memAllocatable.Value() > 0 {
			score += float64(memorySum.Value()) / float64(memAllocatable.Value())
		}

		podInfos = append(podInfos, &podInfo{
			score: score,
			pod:   pod,
		})
	}
	return podInfos
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

import (
	"sync"
)

// podMetadata contains metadata about a single pod
type podMetadata struct {
	stateless bool
}

// PodsMetadata store metadata about pods.
type PodsMetadata struct {
	name     string
	lock     sync.Mutex
	metadata map[string]*podMetadata
}

// NewPodsMetadata created new PodsMetadata instance.
func NewPodsMetadata(name string) *PodsMetadata {
	return &PodsMetadata{
		name:     name,
		metadata: make(map[string]*podMetadata),
	}
}

// SetStaless marks a given pod as stateless.
func (o *PodsMetadata) SetStateless(key string, stateless bool) {
	o.lock.Lock()
	defer o.lock.Unlock()
	if meta, ok := o.metadata[key]; ok {
		meta.stateless = stateless
		return
	}
	o.metadata[key] = &podMetadata{
		stateless: stateless,
	}
}

// FilterStateless returns true iff a pod associated with a given key is
// marked as stateless.
func (o *PodsMetadata) FilterStateless(key string) bool {
	o.lock.Lock()
	defer o.lock.Unlock()
	meta, ok := o.metadata[key]
	return ok && meta.stateless
}
<commit_msg>Fix golint issue caused by typo<commit_after>/*
Copyright 2020 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

import (
	"sync"
)

// podMetadata contains metadata about a single pod
type podMetadata struct {
	stateless bool
}

// PodsMetadata store metadata about pods.
type PodsMetadata struct {
	name     string
	lock     sync.Mutex
	metadata map[string]*podMetadata
}

// NewPodsMetadata created new PodsMetadata instance.
func NewPodsMetadata(name string) *PodsMetadata {
	return &PodsMetadata{
		name:     name,
		metadata: make(map[string]*podMetadata),
	}
}

// SetStateless marks a given pod as stateless.
func (o *PodsMetadata) SetStateless(key string, stateless bool) {
	o.lock.Lock()
	defer o.lock.Unlock()
	if meta, ok := o.metadata[key]; ok {
		meta.stateless = stateless
		return
	}
	o.metadata[key] = &podMetadata{
		stateless: stateless,
	}
}

// FilterStateless returns true iff a pod associated with a given key is
// marked as stateless.
func (o *PodsMetadata) FilterStateless(key string) bool {
	o.lock.Lock()
	defer o.lock.Unlock()
	meta, ok := o.metadata[key]
	return ok && meta.stateless
}
<|endoftext|>


<commit_before>package shutdown

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
)

var h = newHandler()

type handler struct {
	active atomic.Value
	mtx    sync.Mutex
	stack  []func()
}

func newHandler() *handler {
	h := &handler{}
	h.active.Store(false)
	go h.wait()
	return h
}

func IsActive() bool {
	return h.active.Load().(bool)
}

func BeforeExit(f func()) {
	h.BeforeExit(f)
}

func (h *handler) BeforeExit(f func()) {
	h.mtx.Lock()
	h.stack = append(h.stack, f)
	h.mtx.Unlock()
}

func Exit() {
	h.exit(nil, 0)
}

func ExitWithCode(code int) {
	h.exit(nil, code)
}

func Fatal(v ...interface{}) {
	h.Fatal(v)
}

func (h *handler) Fatal(v ...interface{}) {
	h.exit(errors.New(fmt.Sprint(v...)), 1)
}

func (h *handler) wait() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt, os.Signal(syscall.SIGTERM))
	<-ch
	h.exit(nil, 0)
}

func (h *handler) exit(err error, code int) {
	h.mtx.Lock()
	h.active.Store(true)
	for i := len(h.stack) - 1; i >= 0; i-- {
		h.stack[i]()
	}
	if err != nil {
		log.New(os.Stderr, "", log.Lshortfile|log.Lmicroseconds).Output(3, err.Error())
	}
	os.Exit(code)
}
<commit_msg>pkg/shutdown: If panics, shutdown passes them.<commit_after>package shutdown

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
)

var h = newHandler()

type handler struct {
	active atomic.Value
	mtx    sync.Mutex
	stack  []func()
}

func newHandler() *handler {
	h := &handler{}
	h.active.Store(false)
	go h.wait()
	return h
}

func IsActive() bool {
	return h.active.Load().(bool)
}

func BeforeExit(f func()) {
	h.BeforeExit(f)
}

func (h *handler) BeforeExit(f func()) {
	h.mtx.Lock()
	h.stack = append(h.stack, f)
	h.mtx.Unlock()
}

func Exit() {
	h.exit(nil, 0, recover())
}

func ExitWithCode(code int) {
	h.exit(nil, code, recover())
}

func Fatal(v ...interface{}) {
	h.Fatal(v)
}

func (h *handler) Fatal(v ...interface{}) {
	h.exit(errors.New(fmt.Sprint(v...)), 1, recover())
}

func (h *handler) wait() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt, os.Signal(syscall.SIGTERM))
	<-ch
	h.exit(nil, 0, nil)
}

func (h *handler) exit(err error, code int, serious interface{}) {
	h.mtx.Lock()
	h.active.Store(true)
	for i := len(h.stack) - 1; i >= 0; i-- {
		h.stack[i]()
	}
	if serious != nil {
		panic(serious)
	}
	if err != nil {
		log.New(os.Stderr, "", log.Lshortfile|log.Lmicroseconds).Output(3, err.Error())
	}
	os.Exit(code)
}
<|endoftext|>


<commit_before>package config

import (
	"os"
	"testing"
	"time"
)

var (
	jsonFileName string = time.Now().Format("200601020304.json")
	jsonProvider Provider
)

type TestData struct {
	Date time.Time
	Name string
}

func TestJsonSet(t *testing.T) {
	testData := &TestData{
		Date: time.Now(),
		Name: "default",
	}

	// new provider
	jsonProvider, err := New(jsonFileName, "json")
	if err != nil {
		// t.Logf("[JSON New]%v", err.Error())
		t.Fatalf("[JSON New]%v", err.Error())
		return
	}

	n, err := jsonProvider.Set(*testData)
	if err != nil {
		t.Fatalf("[JSON]write %v bytes, error: %v\n", n, err)
	}
}
<commit_msg>to clean the case files.remove json provider test case file.<commit_after>package config

import (
	"os"
	"testing"
	"time"
)

var (
	jsonFileName string = time.Now().Format("200601020304.json")
	jsonProvider Provider
)

type TestData struct {
	Date time.Time
	Name string
}

func TestJsonSet(t *testing.T) {
	testData := &TestData{
		Date: time.Now(),
		Name: "default",
	}

	// new provider
	jsonProvider, err := New(jsonFileName, "json")
	if err != nil {
		// t.Logf("[JSON New]%v", err.Error())
		t.Fatalf("[JSON New]%v", err.Error())
		return
	}

	n, err := jsonProvider.Set(*testData)
	if err != nil {
		t.Fatalf("[JSON]write %v bytes, error: %v\n", n, err)
	}
}

func TestCleanCase(t *testing.T) {
	t.Logf("Remove json file: %s\n", jsonFileName)
	os.Remove(jsonFileName)
}
<|endoftext|>


<commit_before>package slinga

import (
	"github.com/golang/glog"
	"k8s.io/helm/pkg/helm"
	"strings"
	yaml "gopkg.in/yaml.v2"
)

// HelmCodeExecutor is an executor that uses Helm for deployment of apps on kubernetes
type HelmCodeExecutor struct {
	Code *Code
}

func HelmName(str string) string {
	r := strings.NewReplacer("#", "-", "_", "-")
	return r.Replace(str)
}

const (
	//tillerHost = "tiller-deploy.kube-system.svc.cluster.local:44134"
	tillerHost = "kapp-demo-1:40666"
)

func newHelmClient() *helm.Client {
	return helm.NewClient(helm.Host(tillerHost))
}

func findHelmRelease(helmClient *helm.Client, name string) (bool, error) {
	resp, err := helmClient.ListReleases()
	if err != nil {
		return false, nil
	}

	for _, rel := range resp.Releases {
		if rel.Name == name {
			return true, nil
		}
	}

	return false, nil
}

// Install for HelmCodeExecutor runs "helm install" for the corresponding helm chart
func (executor HelmCodeExecutor) Install(key string, codeMetadata map[string]string, codeParams interface{}) error {
	uid := strings.ToLower(HelmName(key))

	chartName := codeMetadata["chartName"]

	helmClient := newHelmClient()

	// TODO check err separately
	if exists, err := findHelmRelease(helmClient, uid); exists && err == nil {
		// TODO log that it's already installed
		// TODO update release just in case
		return nil
	}

	chartPath := GetAptomiPolicyDir() + "/charts/" + chartName + ".tgz"

	vals, err := yaml.Marshal(codeParams)
	if err != nil {
		return err
	}

	glog.Infof("Installing new Helm release '%s' of '%s' (path: %s) with params:\n%s", uid, chartName, chartPath, string(vals))

	// TODO is it good to reuse name?
	_ /*resp*/, err = helmClient.InstallRelease(chartPath, "aptomi", helm.ReleaseName(uid), helm.ValueOverrides(vals), helm.InstallReuseName(true))
	if err != nil {
		return err
	}
	return nil
}

// Update for HelmCodeExecutor runs "helm update" for the corresponding helm chart
func (executor HelmCodeExecutor) Update(key string, labels LabelSet) error {
	// TODO: implement update method
	return nil
}

// Destroy for HelmCodeExecutor runs "helm delete" for the corresponding helm chart
func (executor HelmCodeExecutor) Destroy(key string) error {
	uid := strings.ToLower(HelmName(key))

	helmClient := newHelmClient()

	glog.Infof("Deleting Helm release '%s'", uid)

	if _, err := helmClient.DeleteRelease(uid, helm.DeletePurge(true)); err != nil {
		return err
	}

	return nil
}
<commit_msg>Small Helm exec code improvements<commit_after>package slinga

import (
	"github.com/golang/glog"
	"k8s.io/helm/pkg/helm"
	"strings"
	yaml "gopkg.in/yaml.v2"
)

// HelmCodeExecutor is an executor that uses Helm for deployment of apps on kubernetes
type HelmCodeExecutor struct {
	Code *Code
}

func HelmName(str string) string {
	r := strings.NewReplacer("#", "-", "_", "-")
	return r.Replace(str)
}

const (
	//tillerHost = "tiller-deploy.kube-system.svc.cluster.local:44134"
	tillerHost = "kapp-demo-1:40666"
)

func newHelmClient() *helm.Client {
	return helm.NewClient(helm.Host(tillerHost))
}

func findHelmRelease(helmClient *helm.Client, name string) (bool, error) {
	resp, err := helmClient.ListReleases()
	if err != nil {
		return false, nil
	}

	for _, rel := range resp.Releases {
		if rel.Name == name {
			return true, nil
		}
	}

	return false, nil
}

// Install for HelmCodeExecutor runs "helm install" for the corresponding helm chart
func (executor HelmCodeExecutor) Install(key string, codeMetadata map[string]string, codeParams interface{}) error {
	releaseName := strings.ToLower(HelmName(key))

	chartName := codeMetadata["chartName"]

	helmClient := newHelmClient()

	// TODO check err separately
	if exists, err := findHelmRelease(helmClient, releaseName); exists && err == nil {
		// TODO log that it's already installed
		// TODO update release just in case
		return nil
	}

	chartPath := GetAptomiPolicyDir() + "/charts/" + chartName + ".tgz"

	vals, err := yaml.Marshal(codeParams)
	if err != nil {
		return err
	}

	glog.Infof("Installing new Helm release '%s' of '%s' (path: %s) with params:\n%s", releaseName, chartName, chartPath, string(vals))

	// TODO is it good to reuse name?
	_ /*resp*/, err = helmClient.InstallRelease(chartPath, "aptomi", helm.ReleaseName(releaseName), helm.ValueOverrides(vals), helm.InstallReuseName(true))
	if err != nil {
		return err
	}
	return nil
}

// Update for HelmCodeExecutor runs "helm update" for the corresponding helm chart
func (executor HelmCodeExecutor) Update(key string, labels LabelSet) error {
	// TODO: implement update method
	return nil
}

// Destroy for HelmCodeExecutor runs "helm delete" for the corresponding helm chart
func (executor HelmCodeExecutor) Destroy(key string) error {
	releaseName := strings.ToLower(HelmName(key))

	helmClient := newHelmClient()

	glog.Infof("Deleting Helm release '%s'", releaseName)

	if _, err := helmClient.DeleteRelease(releaseName, helm.DeletePurge(true)); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package config_test

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

	"github.com/fly/config"
)

// In this example our organization is named "podhub", and our project
// namespace is "canary".
//
// In this example we have a file located at /Users/jchen/.config/podhub/canary/config.yaml,
// with the following contents:
//  example:
//    - "a"
//    - "b"
//    - "c"
func ExampleNamespace() {
	type Config struct {
		Example []string `yaml:"example"`
	}

	var err error
	var cfg Config
	var path string
	var cfgNS = config.Namespace{
		Organization: "podhub",
		System:       "canary",
	}

	path, _ = cfgNS.Path()
	fmt.Println("Path to config " + path)

	err = cfgNS.Load(&cfg)
	if err != nil {
		log.Fatalln(err)
	}
	fmt.Println("Contents of cfg " + fmt.Sprint(cfg))
}

func TestExpandUser(t *testing.T) {
	var homeDir string
	var correctPath = homeDir + "/.config/fly/config/config.yaml"
	var err error
	var path string

	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}

	path, err = config.ExpandUser("~/.config/fly/config/config.yaml")

	if err != nil {
		t.Error("Got an error: ", err, ", expecting nil")
	}

	// docs say not to trust /home/travis to be homedir. We'll need to
	// revisit this later.
	if path != correctPath {
		t.Error("Expected ", correctPath, ", got ", path)
	}
}

func TestLoad(t *testing.T) {
	const correctDir = "/etc/fly/config/"
	const correctPath = correctDir + "config.yaml"
	type configExample struct {
		Location string `yaml:"location"`
		Burritos bool   `yaml:"burritos"`
	}

	var correctCfgText = `location: Señor Sisig
burritos: true`
	var correctCfg = configExample{
		Location: "Señor Sisig",
		Burritos: true,
	}
	var err error
	var cfg configExample
	var homeDir string
	var dirMode os.FileMode = 0755
	var fileMode os.FileMode = 0644

	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}

	// Setup
	os.RemoveAll(homeDir + "/.config/fly/config/config.yaml")
	os.MkdirAll(correctDir, dirMode)
	err = ioutil.WriteFile(correctPath, []byte(correctCfgText), fileMode)
	if err != nil {
		t.Error("Got an error writing ", correctPath, ", got an error: ", err)
	}

	// Test
	err = config.Load(correctPath, &cfg)
	if err != nil {
		t.Error("Got an error: ", err, ", expecting nil")
	}

	if cfg != correctCfg {
		t.Error("Expecting ", correctCfg, ", got ", cfg)
	}

	// Teardown
	err = os.RemoveAll(correctPath)
	if err != nil {
		t.Error("Unable to remove file ", correctPath,
			" in teardown, got an error: ", err)
	}
}

func TestUserBase(t *testing.T) {
	const correctUserBase = "~/.config/"
	var userBase string
	userBase = config.UserBase

	if userBase != correctUserBase {
		t.Error("Expecting ", correctUserBase, ", got ", userBase)
	}
}

func TestSystemBase(t *testing.T) {
	const correctSystemBase = "/etc/"
	var systemBase string
	systemBase = config.SystemBase

	if systemBase != correctSystemBase {
		t.Error("Expecting ", correctSystemBase, ", got ", systemBase)
	}
}

func TestNamespacePath(t *testing.T) {
	const correctDir = "/etc/fly/config/"
	const correctPath = correctDir + "config.yaml"
	var cfgNS = config.Namespace{
		Organization: "fly",
		System:       "config",
	}
	var err error
	var path string
	var homeDir string
	var dirMode os.FileMode = 0755
	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}

	// Setup
	os.RemoveAll(homeDir + "/.config/fly/config/config.yaml")
	os.MkdirAll(correctDir, dirMode)
	_, err = os.Create(correctPath)
	if err != nil {
		t.Error("Unable to create file ", correctPath, ", got an error: ", err)
	}

	// Test
	path, err = cfgNS.Path()
	if err != nil {
		t.Error("Got an error: ", err, " expecting nil")
	}
	if path != correctPath {
		t.Error("Expecting ", correctPath, ", got ", path)
	}

	// Teardown
	err = os.RemoveAll(correctPath)
	if err != nil {
		t.Error("Unable to remove file ", correctPath,
			" in teardown, got an error: ", err)
	}
}

func TestNamespaceLoad(t *testing.T) {
	const correctDir = "/etc/fly/config/"
	const correctPath = correctDir + "config.yaml"
	type configExample struct {
		Location string
		Burritos bool
	}

	var correctCfgText = `location: Señor Sisig
burritos: true`
	var correctCfg = configExample{
		Location: "Señor Sisig",
		Burritos: true,
	}
	var cfgNS = config.Namespace{
		Organization: "fly",
		System:       "config",
	}
	var err error
	var cfg configExample
	var homeDir string
	var dirMode os.FileMode = 0755
	var fileMode os.FileMode = 0644

	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}

	// Setup
	os.RemoveAll(homeDir + "/.config/fly/config/config.yaml")
	os.MkdirAll(correctDir, dirMode)
	err = ioutil.WriteFile(correctPath, []byte(correctCfgText), fileMode)
	if err != nil {
		t.Error("Got an error writing ", correctPath, ", got an error: ", err)
	}

	// Test
	err = cfgNS.Load(&cfg)
	if err != nil {
		t.Error("Got an error: ", err, ", expecting nil")
	}

	if cfg != correctCfg {
		t.Error("Expecting ", correctCfg, ", got ", cfg)
	}

	// Teardown
	err = os.RemoveAll(correctPath)
	if err != nil {
		t.Error("Unable to remove file ", correctPath,
			" in teardown, got an error: ", err)
	}
}
<commit_msg>fix correctPath definition<commit_after>package config_test

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

	"github.com/fly/config"
)

// In this example our organization is named "podhub", and our project
// namespace is "canary".
//
// In this example we have a file located at /Users/jchen/.config/podhub/canary/config.yaml,
// with the following contents:
//  example:
//    - "a"
//    - "b"
//    - "c"
func ExampleNamespace() {
	type Config struct {
		Example []string `yaml:"example"`
	}

	var err error
	var cfg Config
	var path string
	var cfgNS = config.Namespace{
		Organization: "podhub",
		System:       "canary",
	}

	path, _ = cfgNS.Path()
	fmt.Println("Path to config " + path)

	err = cfgNS.Load(&cfg)
	if err != nil {
		log.Fatalln(err)
	}
	fmt.Println("Contents of cfg " + fmt.Sprint(cfg))
}

func TestExpandUser(t *testing.T) {
	var homeDir string
	var err error
	var path string

	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}
	var correctPath = homeDir + "/.config/fly/config/config.yaml"

	path, err = config.ExpandUser("~/.config/fly/config/config.yaml")

	if err != nil {
		t.Error("Got an error: ", err, ", expecting nil")
	}

	// docs say not to trust /home/travis to be homedir. We'll need to
	// revisit this later.
	if path != correctPath {
		t.Error("Expected ", correctPath, ", got ", path)
	}
}

func TestLoad(t *testing.T) {
	const correctDir = "/etc/fly/config/"
	type configExample struct {
		Location string `yaml:"location"`
		Burritos bool   `yaml:"burritos"`
	}

	var correctCfgText = `location: Señor Sisig
burritos: true`
	var correctCfg = configExample{
		Location: "Señor Sisig",
		Burritos: true,
	}
	var err error
	var cfg configExample
	var homeDir string
	var dirMode os.FileMode = 0755
	var fileMode os.FileMode = 0644

	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}
	const correctPath = correctDir + "config.yaml"

	// Setup
	os.RemoveAll(homeDir + "/.config/fly/config/config.yaml")
	os.MkdirAll(correctDir, dirMode)
	err = ioutil.WriteFile(correctPath, []byte(correctCfgText), fileMode)
	if err != nil {
		t.Error("Got an error writing ", correctPath, ", got an error: ", err)
	}

	// Test
	err = config.Load(correctPath, &cfg)
	if err != nil {
		t.Error("Got an error: ", err, ", expecting nil")
	}

	if cfg != correctCfg {
		t.Error("Expecting ", correctCfg, ", got ", cfg)
	}

	// Teardown
	err = os.RemoveAll(correctPath)
	if err != nil {
		t.Error("Unable to remove file ", correctPath,
			" in teardown, got an error: ", err)
	}
}

func TestUserBase(t *testing.T) {
	const correctUserBase = "~/.config/"
	var userBase string
	userBase = config.UserBase

	if userBase != correctUserBase {
		t.Error("Expecting ", correctUserBase, ", got ", userBase)
	}
}

func TestSystemBase(t *testing.T) {
	const correctSystemBase = "/etc/"
	var systemBase string
	systemBase = config.SystemBase

	if systemBase != correctSystemBase {
		t.Error("Expecting ", correctSystemBase, ", got ", systemBase)
	}
}

func TestNamespacePath(t *testing.T) {
	const correctDir = "/etc/fly/config/"
	var cfgNS = config.Namespace{
		Organization: "fly",
		System:       "config",
	}
	var err error
	var path string
	var homeDir string
	var dirMode os.FileMode = 0755
	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}
	const correctPath = correctDir + "config.yaml"

	// Setup
	os.RemoveAll(homeDir + "/.config/fly/config/config.yaml")
	os.MkdirAll(correctDir, dirMode)
	_, err = os.Create(correctPath)
	if err != nil {
		t.Error("Unable to create file ", correctPath, ", got an error: ", err)
	}

	// Test
	path, err = cfgNS.Path()
	if err != nil {
		t.Error("Got an error: ", err, " expecting nil")
	}
	if path != correctPath {
		t.Error("Expecting ", correctPath, ", got ", path)
	}

	// Teardown
	err = os.RemoveAll(correctPath)
	if err != nil {
		t.Error("Unable to remove file ", correctPath,
			" in teardown, got an error: ", err)
	}
}

func TestNamespaceLoad(t *testing.T) {
	const correctDir = "/etc/fly/config/"
	type configExample struct {
		Location string
		Burritos bool
	}

	var correctCfgText = `location: Señor Sisig
burritos: true`
	var correctCfg = configExample{
		Location: "Señor Sisig",
		Burritos: true,
	}
	var cfgNS = config.Namespace{
		Organization: "fly",
		System:       "config",
	}
	var err error
	var cfg configExample
	var homeDir string
	var dirMode os.FileMode = 0755
	var fileMode os.FileMode = 0644

	if os.Getenv("TRAVIS") == "true" {
		homeDir = "/home/travis"
	} else {
		homeDir = os.Getenv("HOME")
	}
	const correctPath = correctDir + "config.yaml"

	// Setup
	os.RemoveAll(homeDir + "/.config/fly/config/config.yaml")
	os.MkdirAll(correctDir, dirMode)
	err = ioutil.WriteFile(correctPath, []byte(correctCfgText), fileMode)
	if err != nil {
		t.Error("Got an error writing ", correctPath, ", got an error: ", err)
	}

	// Test
	err = cfgNS.Load(&cfg)
	if err != nil {
		t.Error("Got an error: ", err, ", expecting nil")
	}

	if cfg != correctCfg {
		t.Error("Expecting ", correctCfg, ", got ", cfg)
	}

	// Teardown
	err = os.RemoveAll(correctPath)
	if err != nil {
		t.Error("Unable to remove file ", correctPath,
			" in teardown, got an error: ", err)
	}
}
<|endoftext|>


<commit_before>// +build darwin freebsd openbsd

package term

import (
	"unsafe"

	"golang.org/x/sys/unix"
)

const (
	getTermios = unix.TIOCGETA
	setTermios = unix.TIOCSETA
)

// Termios is the Unix API for terminal I/O.
type Termios unix.Termios

// MakeRaw put the terminal connected to the given file descriptor into raw
// mode and returns the previous state of the terminal so that it can be
// restored.
func MakeRaw(fd uintptr) (*State, error) {
	var oldState State
	if _, _, err := unix.Syscall(unix.SYS_IOCTL, fd, getTermios, uintptr(unsafe.Pointer(&oldState.termios))); err != 0 {
		return nil, err
	}

	newState := oldState.termios
	newState.Iflag &^= (unix.IGNBRK | unix.BRKINT | unix.PARMRK | unix.ISTRIP | unix.INLCR | unix.IGNCR | unix.ICRNL | unix.IXON)
	newState.Oflag &^= unix.OPOST
	newState.Lflag &^= (unix.ECHO | unix.ECHONL | unix.ICANON | unix.ISIG | unix.IEXTEN)
	newState.Cflag &^= (unix.CSIZE | unix.PARENB)
	newState.Cflag |= unix.CS8
	newState.Cc[unix.VMIN] = 1
	newState.Cc[unix.VTIME] = 0

	if _, _, err := unix.Syscall(unix.SYS_IOCTL, fd, setTermios, uintptr(unsafe.Pointer(&newState))); err != 0 {
		return nil, err
	}

	return &oldState, nil
}
<commit_msg>Set OPOST on bsd<commit_after>// +build darwin freebsd openbsd

package term

import (
	"unsafe"

	"golang.org/x/sys/unix"
)

const (
	getTermios = unix.TIOCGETA
	setTermios = unix.TIOCSETA
)

// Termios is the Unix API for terminal I/O.
type Termios unix.Termios

// MakeRaw put the terminal connected to the given file descriptor into raw
// mode and returns the previous state of the terminal so that it can be
// restored.
func MakeRaw(fd uintptr) (*State, error) {
	var oldState State
	if _, _, err := unix.Syscall(unix.SYS_IOCTL, fd, getTermios, uintptr(unsafe.Pointer(&oldState.termios))); err != 0 {
		return nil, err
	}

	newState := oldState.termios
	newState.Iflag &^= (unix.IGNBRK | unix.BRKINT | unix.PARMRK | unix.ISTRIP | unix.INLCR | unix.IGNCR | unix.ICRNL | unix.IXON)
	newState.Oflag |= unix.OPOST
	newState.Lflag &^= (unix.ECHO | unix.ECHONL | unix.ICANON | unix.ISIG | unix.IEXTEN)
	newState.Cflag &^= (unix.CSIZE | unix.PARENB)
	newState.Cflag |= unix.CS8
	newState.Cc[unix.VMIN] = 1
	newState.Cc[unix.VTIME] = 0

	if _, _, err := unix.Syscall(unix.SYS_IOCTL, fd, setTermios, uintptr(unsafe.Pointer(&newState))); err != 0 {
		return nil, err
	}

	return &oldState, nil
}
<|endoftext|>


<commit_before>// +build !windows

package main

import (
	"bytes"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

func configFile() (string, error) {
	dir, err := homeDir()
	if err != nil {
		return "", err
	}

	return filepath.Join(dir, ".terraformrc"), nil
}

func configDir() (string, error) {
	dir, err := homeDir()
	if err != nil {
		return "", err
	}

	return filepath.Join(dir, ".terraform.d"), nil
}

func homeDir() (string, error) {
	// First prefer the HOME environmental variable
	if home := os.Getenv("HOME"); home != "" {
		// FIXME: homeDir gets called from globalPluginDirs during init, before
		// the logging is setup.  We should move meta initializtion outside of
		// init, but in the meantime we just need to silence this output.
		//log.Printf("[DEBUG] Detected home directory from env var: %s", home)

		return home, nil
	}

	// If that fails, try the shell
	var stdout bytes.Buffer
	cmd := exec.Command("sh", "-c", "eval echo ~$USER")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return "", err
	}

	result := strings.TrimSpace(stdout.String())
	if result == "" {
		return "", errors.New("blank output")
	}

	return result, nil
}
<commit_msg>Use build-in method to get user homedir instead of eval on sh<commit_after>// +build !windows

package main

import (
	"errors"
	"os"
	"os/user"
	"path/filepath"
)

func configFile() (string, error) {
	dir, err := homeDir()
	if err != nil {
		return "", err
	}

	return filepath.Join(dir, ".terraformrc"), nil
}

func configDir() (string, error) {
	dir, err := homeDir()
	if err != nil {
		return "", err
	}

	return filepath.Join(dir, ".terraform.d"), nil
}

func homeDir() (string, error) {
	// First prefer the HOME environmental variable
	if home := os.Getenv("HOME"); home != "" {
		// FIXME: homeDir gets called from globalPluginDirs during init, before
		// the logging is setup.  We should move meta initializtion outside of
		// init, but in the meantime we just need to silence this output.
		//log.Printf("[DEBUG] Detected home directory from env var: %s", home)

		return home, nil
	}

	// If that fails, try build-in module
	user, err := user.Current()
	if err != nil {
		return "", err
	}

	if user.HomeDir == "" {
		return "", errors.New("blank output")
	}

	return user.HomeDir, nil
}
<|endoftext|>


<commit_before>package drivers

import (
	"errors"
	"fmt"
	"io"
	"time"

	"github.com/golang/protobuf/ptypes"
	hclog "github.com/hashicorp/go-hclog"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/plugins/base"
	"github.com/hashicorp/nomad/plugins/drivers/proto"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"golang.org/x/net/context"
)

var _ DriverPlugin = &driverPluginClient{}

type driverPluginClient struct {
	*base.BasePluginClient

	client proto.DriverClient
	logger hclog.Logger
}

func (d *driverPluginClient) TaskConfigSchema() (*hclspec.Spec, error) {
	req := &proto.TaskConfigSchemaRequest{}

	resp, err := d.client.TaskConfigSchema(context.Background(), req)
	if err != nil {
		return nil, err
	}

	return resp.Spec, nil
}

func (d *driverPluginClient) Capabilities() (*Capabilities, error) {
	req := &proto.CapabilitiesRequest{}

	resp, err := d.client.Capabilities(context.Background(), req)
	if err != nil {
		return nil, err
	}

	caps := &Capabilities{}
	if resp.Capabilities != nil {
		caps.SendSignals = resp.Capabilities.SendSignals
		caps.Exec = resp.Capabilities.Exec

		switch resp.Capabilities.FsIsolation {
		case proto.DriverCapabilities_NONE:
			caps.FSIsolation = cstructs.FSIsolationNone
		case proto.DriverCapabilities_CHROOT:
			caps.FSIsolation = cstructs.FSIsolationChroot
		case proto.DriverCapabilities_IMAGE:
			caps.FSIsolation = cstructs.FSIsolationImage
		default:
			caps.FSIsolation = cstructs.FSIsolationNone
		}
	}

	return caps, nil
}

// Fingerprint the driver, return a chan that will be pushed to periodically and on changes to health
func (d *driverPluginClient) Fingerprint(ctx context.Context) (<-chan *Fingerprint, error) {
	req := &proto.FingerprintRequest{}

	stream, err := d.client.Fingerprint(ctx, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan *Fingerprint)
	go d.handleFingerprint(ctx, ch, stream)

	return ch, nil
}

func (d *driverPluginClient) handleFingerprint(ctx context.Context, ch chan *Fingerprint, stream proto.Driver_FingerprintClient) {
	defer close(ch)
	for {
		pb, err := stream.Recv()
		if err == io.EOF {
			return
		}
		if err != nil {
			d.logger.Error("error receiving stream from Fingerprint driver RPC", "error", err)
			select {
			case <-ctx.Done():
			case ch <- &Fingerprint{Err: fmt.Errorf("error from RPC stream: %v", err)}:
			}
			return
		}
		f := &Fingerprint{
			Attributes:        pb.Attributes,
			Health:            healthStateFromProto(pb.Health),
			HealthDescription: pb.HealthDescription,
		}

		select {
		case <-ctx.Done():
			return
		case ch <- f:
		}
	}
}

// RecoverTask does internal state recovery to be able to control the task of
// the given TaskHandle
func (d *driverPluginClient) RecoverTask(h *TaskHandle) error {
	req := &proto.RecoverTaskRequest{Handle: taskHandleToProto(h)}

	_, err := d.client.RecoverTask(context.Background(), req)
	return err
}

// StartTask starts execution of a task with the given TaskConfig. A TaskHandle
// is returned to the caller that can be used to recover state of the task,
// should the driver crash or exit prematurely.
func (d *driverPluginClient) StartTask(c *TaskConfig) (*TaskHandle, *cstructs.DriverNetwork, error) {
	req := &proto.StartTaskRequest{
		Task: taskConfigToProto(c),
	}

	resp, err := d.client.StartTask(context.Background(), req)
	if err != nil {
		return nil, nil, err
	}

	var net *cstructs.DriverNetwork
	if resp.NetworkOverride != nil {
		net = &cstructs.DriverNetwork{
			PortMap:       map[string]int{},
			IP:            resp.NetworkOverride.Addr,
			AutoAdvertise: resp.NetworkOverride.AutoAdvertise,
		}
		for k, v := range resp.NetworkOverride.PortMap {
			net.PortMap[k] = int(v)
		}
	}

	return taskHandleFromProto(resp.Handle), net, nil
}

// WaitTask returns a channel that will have an ExitResult pushed to it once when the task
// exits on its own or is killed. If WaitTask is called after the task has exited, the channel
// will immedialy return the ExitResult. WaitTask can be called multiple times for
// the same task without issue.
func (d *driverPluginClient) WaitTask(ctx context.Context, id string) (<-chan *ExitResult, error) {
	ch := make(chan *ExitResult)
	go d.handleWaitTask(ctx, id, ch)
	return ch, nil
}

func (d *driverPluginClient) handleWaitTask(ctx context.Context, id string, ch chan *ExitResult) {
	defer close(ch)
	var result ExitResult
	req := &proto.WaitTaskRequest{
		TaskId: id,
	}

	resp, err := d.client.WaitTask(ctx, req)
	if err != nil {
		result.Err = err
	} else {
		result.ExitCode = int(resp.Result.ExitCode)
		result.Signal = int(resp.Result.Signal)
		result.OOMKilled = resp.Result.OomKilled
		if len(resp.Err) > 0 {
			result.Err = errors.New(resp.Err)
		}
	}
	ch <- &result
}

// StopTask stops the task with the given taskID. A timeout and signal can be
// given to control a graceful termination of the task. The driver will send the
// given signal to the task and wait for the given timeout for it to exit. If the
// task does not exit within the timeout it will be forcefully killed.
func (d *driverPluginClient) StopTask(taskID string, timeout time.Duration, signal string) error {
	req := &proto.StopTaskRequest{
		TaskId:  taskID,
		Timeout: ptypes.DurationProto(timeout),
		Signal:  signal,
	}

	_, err := d.client.StopTask(context.Background(), req)
	return err
}

// DestroyTask removes the task from the driver's in memory state. The task
// cannot be running unless force is set to true. If force is set to true the
// driver will forcefully terminate the task before removing it.
func (d *driverPluginClient) DestroyTask(taskID string, force bool) error {
	req := &proto.DestroyTaskRequest{
		TaskId: taskID,
		Force:  force,
	}

	_, err := d.client.DestroyTask(context.Background(), req)
	return err
}

// InspectTask returns status information for a task
func (d *driverPluginClient) InspectTask(taskID string) (*TaskStatus, error) {
	req := &proto.InspectTaskRequest{TaskId: taskID}

	resp, err := d.client.InspectTask(context.Background(), req)
	if err != nil {
		return nil, err
	}

	status, err := taskStatusFromProto(resp.Task)
	if err != nil {
		return nil, err
	}

	if resp.Driver != nil {
		status.DriverAttributes = resp.Driver.Attributes
	}
	if resp.NetworkOverride != nil {
		status.NetworkOverride = &cstructs.DriverNetwork{
			PortMap:       map[string]int{},
			IP:            resp.NetworkOverride.Addr,
			AutoAdvertise: resp.NetworkOverride.AutoAdvertise,
		}
		for k, v := range resp.NetworkOverride.PortMap {
			status.NetworkOverride.PortMap[k] = int(v)
		}
	}

	return status, nil
}

// TaskStats returns resource usage statistics for the task
func (d *driverPluginClient) TaskStats(taskID string) (*cstructs.TaskResourceUsage, error) {
	req := &proto.TaskStatsRequest{TaskId: taskID}

	resp, err := d.client.TaskStats(context.Background(), req)
	if err != nil {
		return nil, err
	}

	stats, err := taskStatsFromProto(resp.Stats)
	if err != nil {
		return nil, err
	}

	return stats, nil
}

// TaskEvents returns a channel that will receive events from the driver about all
// tasks such as lifecycle events, terminal errors, etc.
func (d *driverPluginClient) TaskEvents(ctx context.Context) (<-chan *TaskEvent, error) {
	req := &proto.TaskEventsRequest{}
	stream, err := d.client.TaskEvents(ctx, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan *TaskEvent)
	go d.handleTaskEvents(ch, stream)
	return ch, nil
}

func (d *driverPluginClient) handleTaskEvents(ch chan *TaskEvent, stream proto.Driver_TaskEventsClient) {
	defer close(ch)
	for {
		ev, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			d.logger.Error("error receiving stream from TaskEvents driver RPC", "error", err)
			ch <- &TaskEvent{Err: err}
			break
		}
		timestamp, _ := ptypes.Timestamp(ev.Timestamp)
		event := &TaskEvent{
			TaskID:      ev.TaskId,
			Annotations: ev.Annotations,
			Message:     ev.Message,
			Timestamp:   timestamp,
		}
		ch <- event
	}
}

// SignalTask will send the given signal to the specified task
func (d *driverPluginClient) SignalTask(taskID string, signal string) error {
	req := &proto.SignalTaskRequest{
		TaskId: taskID,
		Signal: signal,
	}
	_, err := d.client.SignalTask(context.Background(), req)
	return err
}

// ExecTask will run the given command within the execution context of the task.
// The driver will wait for the given timeout for the command to complete before
// terminating it. The stdout and stderr of the command will be return to the caller,
// along with other exit information such as exit code.
func (d *driverPluginClient) ExecTask(taskID string, cmd []string, timeout time.Duration) (*ExecTaskResult, error) {
	req := &proto.ExecTaskRequest{
		TaskId:  taskID,
		Command: cmd,
		Timeout: ptypes.DurationProto(timeout),
	}

	resp, err := d.client.ExecTask(context.Background(), req)
	if err != nil {
		return nil, err
	}

	result := &ExecTaskResult{
		Stdout:     resp.Stdout,
		Stderr:     resp.Stderr,
		ExitResult: exitResultFromProto(resp.Result),
	}

	return result, nil

}
<commit_msg>drivers: only log non-cancellation errors<commit_after>package drivers

import (
	"errors"
	"fmt"
	"io"
	"time"

	"github.com/golang/protobuf/ptypes"
	hclog "github.com/hashicorp/go-hclog"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/plugins/base"
	"github.com/hashicorp/nomad/plugins/drivers/proto"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"golang.org/x/net/context"
)

var _ DriverPlugin = &driverPluginClient{}

type driverPluginClient struct {
	*base.BasePluginClient

	client proto.DriverClient
	logger hclog.Logger
}

func (d *driverPluginClient) TaskConfigSchema() (*hclspec.Spec, error) {
	req := &proto.TaskConfigSchemaRequest{}

	resp, err := d.client.TaskConfigSchema(context.Background(), req)
	if err != nil {
		return nil, err
	}

	return resp.Spec, nil
}

func (d *driverPluginClient) Capabilities() (*Capabilities, error) {
	req := &proto.CapabilitiesRequest{}

	resp, err := d.client.Capabilities(context.Background(), req)
	if err != nil {
		return nil, err
	}

	caps := &Capabilities{}
	if resp.Capabilities != nil {
		caps.SendSignals = resp.Capabilities.SendSignals
		caps.Exec = resp.Capabilities.Exec

		switch resp.Capabilities.FsIsolation {
		case proto.DriverCapabilities_NONE:
			caps.FSIsolation = cstructs.FSIsolationNone
		case proto.DriverCapabilities_CHROOT:
			caps.FSIsolation = cstructs.FSIsolationChroot
		case proto.DriverCapabilities_IMAGE:
			caps.FSIsolation = cstructs.FSIsolationImage
		default:
			caps.FSIsolation = cstructs.FSIsolationNone
		}
	}

	return caps, nil
}

// Fingerprint the driver, return a chan that will be pushed to periodically and on changes to health
func (d *driverPluginClient) Fingerprint(ctx context.Context) (<-chan *Fingerprint, error) {
	req := &proto.FingerprintRequest{}

	stream, err := d.client.Fingerprint(ctx, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan *Fingerprint)
	go d.handleFingerprint(ctx, ch, stream)

	return ch, nil
}

func (d *driverPluginClient) handleFingerprint(ctx context.Context, ch chan *Fingerprint, stream proto.Driver_FingerprintClient) {
	defer close(ch)
	for {
		pb, err := stream.Recv()
		if err == io.EOF {
			return
		}
		if err != nil {
			select {
			case <-ctx.Done():
			case ch <- &Fingerprint{Err: fmt.Errorf("error from RPC stream: %v", err)}:
				d.logger.Error("error receiving stream from Fingerprint driver RPC", "error", err)
			}
			return
		}
		f := &Fingerprint{
			Attributes:        pb.Attributes,
			Health:            healthStateFromProto(pb.Health),
			HealthDescription: pb.HealthDescription,
		}

		select {
		case <-ctx.Done():
			return
		case ch <- f:
		}
	}
}

// RecoverTask does internal state recovery to be able to control the task of
// the given TaskHandle
func (d *driverPluginClient) RecoverTask(h *TaskHandle) error {
	req := &proto.RecoverTaskRequest{Handle: taskHandleToProto(h)}

	_, err := d.client.RecoverTask(context.Background(), req)
	return err
}

// StartTask starts execution of a task with the given TaskConfig. A TaskHandle
// is returned to the caller that can be used to recover state of the task,
// should the driver crash or exit prematurely.
func (d *driverPluginClient) StartTask(c *TaskConfig) (*TaskHandle, *cstructs.DriverNetwork, error) {
	req := &proto.StartTaskRequest{
		Task: taskConfigToProto(c),
	}

	resp, err := d.client.StartTask(context.Background(), req)
	if err != nil {
		return nil, nil, err
	}

	var net *cstructs.DriverNetwork
	if resp.NetworkOverride != nil {
		net = &cstructs.DriverNetwork{
			PortMap:       map[string]int{},
			IP:            resp.NetworkOverride.Addr,
			AutoAdvertise: resp.NetworkOverride.AutoAdvertise,
		}
		for k, v := range resp.NetworkOverride.PortMap {
			net.PortMap[k] = int(v)
		}
	}

	return taskHandleFromProto(resp.Handle), net, nil
}

// WaitTask returns a channel that will have an ExitResult pushed to it once when the task
// exits on its own or is killed. If WaitTask is called after the task has exited, the channel
// will immedialy return the ExitResult. WaitTask can be called multiple times for
// the same task without issue.
func (d *driverPluginClient) WaitTask(ctx context.Context, id string) (<-chan *ExitResult, error) {
	ch := make(chan *ExitResult)
	go d.handleWaitTask(ctx, id, ch)
	return ch, nil
}

func (d *driverPluginClient) handleWaitTask(ctx context.Context, id string, ch chan *ExitResult) {
	defer close(ch)
	var result ExitResult
	req := &proto.WaitTaskRequest{
		TaskId: id,
	}

	resp, err := d.client.WaitTask(ctx, req)
	if err != nil {
		result.Err = err
	} else {
		result.ExitCode = int(resp.Result.ExitCode)
		result.Signal = int(resp.Result.Signal)
		result.OOMKilled = resp.Result.OomKilled
		if len(resp.Err) > 0 {
			result.Err = errors.New(resp.Err)
		}
	}
	ch <- &result
}

// StopTask stops the task with the given taskID. A timeout and signal can be
// given to control a graceful termination of the task. The driver will send the
// given signal to the task and wait for the given timeout for it to exit. If the
// task does not exit within the timeout it will be forcefully killed.
func (d *driverPluginClient) StopTask(taskID string, timeout time.Duration, signal string) error {
	req := &proto.StopTaskRequest{
		TaskId:  taskID,
		Timeout: ptypes.DurationProto(timeout),
		Signal:  signal,
	}

	_, err := d.client.StopTask(context.Background(), req)
	return err
}

// DestroyTask removes the task from the driver's in memory state. The task
// cannot be running unless force is set to true. If force is set to true the
// driver will forcefully terminate the task before removing it.
func (d *driverPluginClient) DestroyTask(taskID string, force bool) error {
	req := &proto.DestroyTaskRequest{
		TaskId: taskID,
		Force:  force,
	}

	_, err := d.client.DestroyTask(context.Background(), req)
	return err
}

// InspectTask returns status information for a task
func (d *driverPluginClient) InspectTask(taskID string) (*TaskStatus, error) {
	req := &proto.InspectTaskRequest{TaskId: taskID}

	resp, err := d.client.InspectTask(context.Background(), req)
	if err != nil {
		return nil, err
	}

	status, err := taskStatusFromProto(resp.Task)
	if err != nil {
		return nil, err
	}

	if resp.Driver != nil {
		status.DriverAttributes = resp.Driver.Attributes
	}
	if resp.NetworkOverride != nil {
		status.NetworkOverride = &cstructs.DriverNetwork{
			PortMap:       map[string]int{},
			IP:            resp.NetworkOverride.Addr,
			AutoAdvertise: resp.NetworkOverride.AutoAdvertise,
		}
		for k, v := range resp.NetworkOverride.PortMap {
			status.NetworkOverride.PortMap[k] = int(v)
		}
	}

	return status, nil
}

// TaskStats returns resource usage statistics for the task
func (d *driverPluginClient) TaskStats(taskID string) (*cstructs.TaskResourceUsage, error) {
	req := &proto.TaskStatsRequest{TaskId: taskID}

	resp, err := d.client.TaskStats(context.Background(), req)
	if err != nil {
		return nil, err
	}

	stats, err := taskStatsFromProto(resp.Stats)
	if err != nil {
		return nil, err
	}

	return stats, nil
}

// TaskEvents returns a channel that will receive events from the driver about all
// tasks such as lifecycle events, terminal errors, etc.
func (d *driverPluginClient) TaskEvents(ctx context.Context) (<-chan *TaskEvent, error) {
	req := &proto.TaskEventsRequest{}
	stream, err := d.client.TaskEvents(ctx, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan *TaskEvent)
	go d.handleTaskEvents(ch, stream)
	return ch, nil
}

func (d *driverPluginClient) handleTaskEvents(ch chan *TaskEvent, stream proto.Driver_TaskEventsClient) {
	defer close(ch)
	for {
		ev, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			d.logger.Error("error receiving stream from TaskEvents driver RPC", "error", err)
			ch <- &TaskEvent{Err: err}
			break
		}
		timestamp, _ := ptypes.Timestamp(ev.Timestamp)
		event := &TaskEvent{
			TaskID:      ev.TaskId,
			Annotations: ev.Annotations,
			Message:     ev.Message,
			Timestamp:   timestamp,
		}
		ch <- event
	}
}

// SignalTask will send the given signal to the specified task
func (d *driverPluginClient) SignalTask(taskID string, signal string) error {
	req := &proto.SignalTaskRequest{
		TaskId: taskID,
		Signal: signal,
	}
	_, err := d.client.SignalTask(context.Background(), req)
	return err
}

// ExecTask will run the given command within the execution context of the task.
// The driver will wait for the given timeout for the command to complete before
// terminating it. The stdout and stderr of the command will be return to the caller,
// along with other exit information such as exit code.
func (d *driverPluginClient) ExecTask(taskID string, cmd []string, timeout time.Duration) (*ExecTaskResult, error) {
	req := &proto.ExecTaskRequest{
		TaskId:  taskID,
		Command: cmd,
		Timeout: ptypes.DurationProto(timeout),
	}

	resp, err := d.client.ExecTask(context.Background(), req)
	if err != nil {
		return nil, err
	}

	result := &ExecTaskResult{
		Stdout:     resp.Stdout,
		Stderr:     resp.Stderr,
		ExitResult: exitResultFromProto(resp.Result),
	}

	return result, nil

}
<|endoftext|>


<commit_before>package cpu

import (
	"fmt"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/inputs/system"
	"github.com/shirou/gopsutil/cpu"
)

type CPUStats struct {
	ps        system.PS
	lastStats map[string]cpu.TimesStat

	PerCPU         bool `toml:"percpu"`
	TotalCPU       bool `toml:"totalcpu"`
	CollectCPUTime bool `toml:"collect_cpu_time"`
	ReportActive   bool `toml:"report_active"`
}

func NewCPUStats(ps system.PS) *CPUStats {
	return &CPUStats{
		ps:             ps,
		CollectCPUTime: true,
		ReportActive:   true,
	}
}

func (_ *CPUStats) Description() string {
	return "Read metrics about cpu usage"
}

var sampleConfig = `
  ## Whether to report per-cpu stats or not
  percpu = true
  ## Whether to report total system cpu stats or not
  totalcpu = true
  ## If true, collect raw CPU time metrics.
  collect_cpu_time = false
  ## If true, compute and report the sum of all non-idle CPU states.
  report_active = false
`

func (_ *CPUStats) SampleConfig() string {
	return sampleConfig
}

func (s *CPUStats) Gather(acc telegraf.Accumulator) error {
	times, err := s.ps.CPUTimes(s.PerCPU, s.TotalCPU)
	if err != nil {
		return fmt.Errorf("error getting CPU info: %s", err)
	}
	now := time.Now()

	for _, cts := range times {
		tags := map[string]string{
			"cpu": cts.CPU,
		}

		total := totalCpuTime(cts)
		active := activeCpuTime(cts)

		if s.CollectCPUTime {
			// Add cpu time metrics
			fieldsC := map[string]interface{}{
				"time_user":       cts.User,
				"time_system":     cts.System,
				"time_idle":       cts.Idle,
				"time_nice":       cts.Nice,
				"time_iowait":     cts.Iowait,
				"time_irq":        cts.Irq,
				"time_softirq":    cts.Softirq,
				"time_steal":      cts.Steal,
				"time_guest":      cts.Guest,
				"time_guest_nice": cts.GuestNice,
			}
			if s.ReportActive {
				fieldsC["time_active"] = activeCpuTime(cts)
			}
			acc.AddCounter("cpu", fieldsC, tags, now)
		}

		// Add in percentage
		if len(s.lastStats) == 0 {
			// If it's the 1st gather, can't get CPU Usage stats yet
			continue
		}

		lastCts, ok := s.lastStats[cts.CPU]
		if !ok {
			continue
		}
		lastTotal := totalCpuTime(lastCts)
		lastActive := activeCpuTime(lastCts)
		totalDelta := total - lastTotal

		if totalDelta < 0 {
			err = fmt.Errorf("Error: current total CPU time is less than previous total CPU time")
			break
		}

		if totalDelta == 0 {
			continue
		}

		fieldsG := map[string]interface{}{
			"usage_user":       100 * (cts.User - lastCts.User - (cts.Guest - lastCts.Guest)) / totalDelta,
			"usage_system":     100 * (cts.System - lastCts.System) / totalDelta,
			"usage_idle":       100 * (cts.Idle - lastCts.Idle) / totalDelta,
			"usage_nice":       100 * (cts.Nice - lastCts.Nice - (cts.GuestNice - lastCts.GuestNice)) / totalDelta,
			"usage_iowait":     100 * (cts.Iowait - lastCts.Iowait) / totalDelta,
			"usage_irq":        100 * (cts.Irq - lastCts.Irq) / totalDelta,
			"usage_softirq":    100 * (cts.Softirq - lastCts.Softirq) / totalDelta,
			"usage_steal":      100 * (cts.Steal - lastCts.Steal) / totalDelta,
			"usage_guest":      100 * (cts.Guest - lastCts.Guest) / totalDelta,
			"usage_guest_nice": 100 * (cts.GuestNice - lastCts.GuestNice) / totalDelta,
		}
		if s.ReportActive {
			fieldsG["usage_active"] = 100 * (active - lastActive) / totalDelta
		}
		acc.AddGauge("cpu", fieldsG, tags, now)
	}

	s.lastStats = make(map[string]cpu.TimesStat)
	for _, cts := range times {
		s.lastStats[cts.CPU] = cts
	}

	return err
}

func totalCpuTime(t cpu.TimesStat) float64 {
	total := t.User + t.System + t.Nice + t.Iowait + t.Irq + t.Softirq + t.Steal +
		t.Idle
	return total
}

func activeCpuTime(t cpu.TimesStat) float64 {
	active := totalCpuTime(t) - t.Idle
	return active
}

func init() {
	inputs.Add("cpu", func() telegraf.Input {
		return &CPUStats{
			PerCPU:   true,
			TotalCPU: true,
			ps:       system.NewSystemPS(),
		}
	})
}
<commit_msg>Unify comments style in the CPU input (#8605)<commit_after>package cpu

import (
	"fmt"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/inputs/system"
	"github.com/shirou/gopsutil/cpu"
)

type CPUStats struct {
	ps        system.PS
	lastStats map[string]cpu.TimesStat

	PerCPU         bool `toml:"percpu"`
	TotalCPU       bool `toml:"totalcpu"`
	CollectCPUTime bool `toml:"collect_cpu_time"`
	ReportActive   bool `toml:"report_active"`
}

func NewCPUStats(ps system.PS) *CPUStats {
	return &CPUStats{
		ps:             ps,
		CollectCPUTime: true,
		ReportActive:   true,
	}
}

func (_ *CPUStats) Description() string {
	return "Read metrics about cpu usage"
}

var sampleConfig = `
  ## Whether to report per-cpu stats or not
  percpu = true
  ## Whether to report total system cpu stats or not
  totalcpu = true
  ## If true, collect raw CPU time metrics
  collect_cpu_time = false
  ## If true, compute and report the sum of all non-idle CPU states
  report_active = false
`

func (_ *CPUStats) SampleConfig() string {
	return sampleConfig
}

func (s *CPUStats) Gather(acc telegraf.Accumulator) error {
	times, err := s.ps.CPUTimes(s.PerCPU, s.TotalCPU)
	if err != nil {
		return fmt.Errorf("error getting CPU info: %s", err)
	}
	now := time.Now()

	for _, cts := range times {
		tags := map[string]string{
			"cpu": cts.CPU,
		}

		total := totalCpuTime(cts)
		active := activeCpuTime(cts)

		if s.CollectCPUTime {
			// Add cpu time metrics
			fieldsC := map[string]interface{}{
				"time_user":       cts.User,
				"time_system":     cts.System,
				"time_idle":       cts.Idle,
				"time_nice":       cts.Nice,
				"time_iowait":     cts.Iowait,
				"time_irq":        cts.Irq,
				"time_softirq":    cts.Softirq,
				"time_steal":      cts.Steal,
				"time_guest":      cts.Guest,
				"time_guest_nice": cts.GuestNice,
			}
			if s.ReportActive {
				fieldsC["time_active"] = activeCpuTime(cts)
			}
			acc.AddCounter("cpu", fieldsC, tags, now)
		}

		// Add in percentage
		if len(s.lastStats) == 0 {
			// If it's the 1st gather, can't get CPU Usage stats yet
			continue
		}

		lastCts, ok := s.lastStats[cts.CPU]
		if !ok {
			continue
		}
		lastTotal := totalCpuTime(lastCts)
		lastActive := activeCpuTime(lastCts)
		totalDelta := total - lastTotal

		if totalDelta < 0 {
			err = fmt.Errorf("Error: current total CPU time is less than previous total CPU time")
			break
		}

		if totalDelta == 0 {
			continue
		}

		fieldsG := map[string]interface{}{
			"usage_user":       100 * (cts.User - lastCts.User - (cts.Guest - lastCts.Guest)) / totalDelta,
			"usage_system":     100 * (cts.System - lastCts.System) / totalDelta,
			"usage_idle":       100 * (cts.Idle - lastCts.Idle) / totalDelta,
			"usage_nice":       100 * (cts.Nice - lastCts.Nice - (cts.GuestNice - lastCts.GuestNice)) / totalDelta,
			"usage_iowait":     100 * (cts.Iowait - lastCts.Iowait) / totalDelta,
			"usage_irq":        100 * (cts.Irq - lastCts.Irq) / totalDelta,
			"usage_softirq":    100 * (cts.Softirq - lastCts.Softirq) / totalDelta,
			"usage_steal":      100 * (cts.Steal - lastCts.Steal) / totalDelta,
			"usage_guest":      100 * (cts.Guest - lastCts.Guest) / totalDelta,
			"usage_guest_nice": 100 * (cts.GuestNice - lastCts.GuestNice) / totalDelta,
		}
		if s.ReportActive {
			fieldsG["usage_active"] = 100 * (active - lastActive) / totalDelta
		}
		acc.AddGauge("cpu", fieldsG, tags, now)
	}

	s.lastStats = make(map[string]cpu.TimesStat)
	for _, cts := range times {
		s.lastStats[cts.CPU] = cts
	}

	return err
}

func totalCpuTime(t cpu.TimesStat) float64 {
	total := t.User + t.System + t.Nice + t.Iowait + t.Irq + t.Softirq + t.Steal +
		t.Idle
	return total
}

func activeCpuTime(t cpu.TimesStat) float64 {
	active := totalCpuTime(t) - t.Idle
	return active
}

func init() {
	inputs.Add("cpu", func() telegraf.Input {
		return &CPUStats{
			PerCPU:   true,
			TotalCPU: true,
			ps:       system.NewSystemPS(),
		}
	})
}
<|endoftext|>