<commit_before>// Package ns provides name server methods for selected name server(s)
package ns

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/miekg/dns"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/data"
)

const (
	publicDNSHost      = "http://public-dns.info"
	publicDNSNodesPath = "/nameservers.csv"
)

// A Host represents a name server host
type Host struct {
	IP      string
	Alpha2  string
	Country string
	City    string
}

// A Request represents a name server request
type Request struct {
	Country      string
	City         string
	Target       string
	Type         uint16
	Host         string
	Hosts        []Host
	TraceEnabled bool
}

// NewRequest creates a new dns request object
func NewRequest() *Request {
	return &Request{Host: ""}
}

// SetOptions passes arguments to appropriate variable
func (d *Request) SetOptions(args, prompt string) bool {
	d.Host = ""
	d.TraceEnabled = false
	d.Type = dns.TypeANY

	nArgs, flag := cli.Flag(args)

	// show help
	if _, ok := flag["help"]; ok || len(nArgs) < 1 {
		help()
		return false
	}

	for _, a := range strings.Fields(nArgs) {
		if a[0] == '@' {
			d.Host = a[1:]
			d.City = ""
			continue
		}
		if t, ok := dns.StringToType[strings.ToUpper(a)]; ok {
			d.Type = t
			continue
		}
		if a == "+trace" {
			d.TraceEnabled = true
			continue
		}
		d.Target = a
	}

	p := strings.Split(prompt, "/")

	if d.Host == "" {
		if p[0] == "local" || len(p) < 3 {
			config, _ := dns.ClientConfigFromFile("/etc/resolv.conf")
			d.Host = config.Servers[0]
			d.City = "your local dns server"
		} else {
			d.ChkNode(p[2])
		}
	}
	return true
}

// Init configure dns command and fetch name servers
func (d *Request) Init() {
	if !d.cache("validate") {
		d.Hosts = fetchNSHosts()
		d.cache("write")
	} else {
		d.cache("read")
	}
}

// CountryList init the connect contry items
func (d *Request) CountryList() []string {
	var countries []string
	for _, host := range d.Hosts {
		countries = append(countries, host.Country)
	}
	countries = uniqStrSlice(countries)
	sort.Strings(countries)
	return countries
}

// NodeList gets the node city items
func (d *Request) NodeList() []string {
	var node []string
	for _, host := range d.Hosts {
		if host.Country == d.Country {
			node = append(node, host.City)
		}
	}
	sort.Strings(node)
	return node
}

// ChkCountry validates and set requested country
func (d *Request) ChkCountry(country string) bool {
	country = strings.ToLower(country)
	for _, h := range d.Hosts {
		if country == h.Country {
			d.Country = country
			return true
		}
	}

	return false
}

// ChkNode set requested country
func (d *Request) ChkNode(city string) bool {
	city = strings.ToLower(city)
	for _, h := range d.Hosts {
		if d.Country == h.Country && city == h.City {
			d.Host = h.IP
			d.City = h.City
			return true
		}
	}
	return false
}

// Local set host to nothing means local
func (d *Request) Local() {
	d.Host = ""
	d.Country = ""
}

// Dig looks up name server w/ trace feature
func (d *Request) Dig() {
	if !d.TraceEnabled {
		d.RunDig()
	} else {
		d.RunDigTrace()
	}
}

// RunDig looks up name server
func (d *Request) RunDig() {
	c := new(dns.Client)
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(d.Target), d.Type)
	m.RecursionDesired = true
	m.RecursionAvailable = true

	println("Trying to query server:", d.Host, d.Country, d.City)

	r, rtt, err := c.Exchange(m, d.Host+":53")
	if err != nil {
		println(err.Error())
		return
	}
	// Answer
	println(r.MsgHdr.String())
	for _, a := range r.Answer {
		fmt.Println(a)
	}
	// Extra info
	if len(r.Extra) > 0 {
		println("\n;; ADDITIONAL SECTION:")
		for _, a := range r.Extra {
			fmt.Println(a)
		}
	}
	fmt.Printf(";; Query time: %d ms\n", rtt/1e6)

	// CHAOS
	c.Timeout = ((rtt / 1e6) + 100) * time.Millisecond
	fmt.Printf("\n;; CHAOS CLASS BIND\n")
	for _, q := range []string{"version.bind.", "hostname.bind."} {
		m.Question[0] = dns.Question{q, dns.TypeTXT, dns.ClassCHAOS}
		r, _, err = c.Exchange(m, d.Host+":53")
		if err != nil {
			continue
		}
		for _, a := range r.Answer {
			fmt.Println(a)
		}
	}
}

// RunDigTrace handles dig trace
func (d *Request) RunDigTrace() {
	var (
		nss  = []string{d.Host}
		err  error
		host string
		rtt  time.Duration
		r    *dns.Msg
	)
	c := new(dns.Client)
	m := new(dns.Msg)
	m.RecursionDesired = true
	q := ""

	domain := []string{""}
	domain = append(domain, strings.Split(dns.Fqdn(d.Target), ".")...)
	for i := range domain {
		if i != 1 && i != len(domain)-1 {
			q = domain[len(domain)-i-1] + "." + q
		} else {
			q = domain[len(domain)-i-1] + q
		}

		if i != len(domain)-1 {
			m.SetQuestion(q, dns.TypeNS)
		} else {
			m.SetQuestion(q, d.Type)
		}

		for _, host = range nss {
			r, rtt, err = c.Exchange(m, host+":53")
			if err != nil {
				println(err.Error())
			} else {
				break
			}
		}

		nss = nss[:0]

		for _, a := range r.Answer {
			println(a.String())
			if a.Header().Rrtype == dns.TypeNS {
				nss = append(nss, strings.Fields(a.String())[4])
			}
		}
		for _, a := range r.Ns {
			println(a.String())
			nss = append(nss, strings.Fields(a.String())[4])
		}

		fmt.Printf("from: %s#53 in %d ms\n", host, rtt/1e6)
	}
}

// cache provides caching for name servers
func (d *Request) cache(r string) bool {
	switch r {
	case "read":
		b, err := ioutil.ReadFile("/tmp/mylg.ns")
		if err != nil {
			panic(err.Error())
		}
		d.Hosts = d.Hosts[:0]
		r := bytes.NewBuffer(b)
		s := bufio.NewScanner(r)
		for s.Scan() {
			csv := strings.Split(s.Text(), ";")
			if len(csv) != 4 {
				continue
			}
			d.Hosts = append(d.Hosts, Host{Alpha2: csv[0], Country: csv[1], City: csv[2], IP: csv[3]})
		}
	case "write":
		var data []string
		for _, h := range d.Hosts {
			data = append(data, fmt.Sprintf("%s;%s;%s;%s", h.Alpha2, h.Country, h.City, h.IP))
		}
		err := ioutil.WriteFile("/tmp/mylg.ns", []byte(strings.Join(data, "\n")), 0644)
		if err != nil {
			panic(err.Error())
		}
	case "validate":
		f, err := os.Stat("/tmp/mylg.ns")
		if err != nil {
			return false
		}
		d := time.Since(f.ModTime())
		if d.Hours() > 48 {
			return false
		}
	}
	return true
}

// Fetch name servers from public-dns.info
func fetchNSHosts() []Host {
	var (
		hosts   []Host
		city    string
		counter = make(map[string]int)
		chkDup  = make(map[string]int)
	)
	resp, err := http.Get(publicDNSHost + publicDNSNodesPath)
	if err != nil {
		println(err.Error())
		return []Host{}
	}
	if resp.StatusCode != 200 {
		println("error: public dns is not available")
		return []Host{}
	}
	defer resp.Body.Close()
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		csv := strings.Split(scanner.Text(), ",")
		if csv[3] != "\"\"" {
			if name, ok := data.Country[csv[2]]; ok && counter[csv[2]] < 5 {
				name = strings.ToLower(name)
				city = strings.ToLower(csv[3])
				chkDup[name+city] += 1
				if chkDup[name+city] > 1 {
					city = fmt.Sprintf("%s0%d", city, chkDup[name+city]-1)
				}

				hosts = append(hosts, Host{IP: csv[0], Alpha2: csv[2], Country: name, City: city})
				counter[csv[2]]++
			}
		}
	}
	return hosts
}

// uniqStrSlice return unique slice
func uniqStrSlice(src []string) []string {
	var rst []string
	tmp := make(map[string]struct{})
	for _, s := range src {
		tmp[s] = struct{}{}
	}
	for s := range tmp {
		rst = append(rst, s)
	}
	return rst
}

// help
func help() {
	fmt.Println(`
    usage:
          dig [@local-server] host [options]
    options:
          +trace
    Example:
          dig google.com
          dig @8.8.8.8 yahoo.com
          dig google.com +trace
          dig google.com MX
	`)

}
<commit_msg>resolved #17, fixed ipv6 dns server issue<commit_after>// Package ns provides name server methods for selected name server(s)
package ns

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/miekg/dns"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/data"
)

const (
	publicDNSHost      = "http://public-dns.info"
	publicDNSNodesPath = "/nameservers.csv"
)

// A Host represents a name server host
type Host struct {
	IP      string
	Alpha2  string
	Country string
	City    string
}

// A Request represents a name server request
type Request struct {
	Country      string
	City         string
	Target       string
	Type         uint16
	Host         string
	Hosts        []Host
	TraceEnabled bool
}

// NewRequest creates a new dns request object
func NewRequest() *Request {
	return &Request{Host: ""}
}

// SetOptions passes arguments to appropriate variable
func (d *Request) SetOptions(args, prompt string) bool {
	d.Host = ""
	d.TraceEnabled = false
	d.Type = dns.TypeANY

	nArgs, flag := cli.Flag(args)

	// show help
	if _, ok := flag["help"]; ok || len(nArgs) < 1 {
		help()
		return false
	}

	for _, a := range strings.Fields(nArgs) {
		if a[0] == '@' {
			d.Host = a[1:]
			d.City = ""
			continue
		}
		if t, ok := dns.StringToType[strings.ToUpper(a)]; ok {
			d.Type = t
			continue
		}
		if a == "+trace" {
			d.TraceEnabled = true
			continue
		}
		d.Target = a
	}

	p := strings.Split(prompt, "/")

	if d.Host == "" {
		if p[0] == "local" || len(p) < 3 {
			config, _ := dns.ClientConfigFromFile("/etc/resolv.conf")
			d.Host = config.Servers[0]
			d.City = "your local dns server"
		} else {
			d.ChkNode(p[2])
		}
	}
	return true
}

// Init configure dns command and fetch name servers
func (d *Request) Init() {
	if !d.cache("validate") {
		d.Hosts = fetchNSHosts()
		d.cache("write")
	} else {
		d.cache("read")
	}
}

// CountryList init the connect contry items
func (d *Request) CountryList() []string {
	var countries []string
	for _, host := range d.Hosts {
		countries = append(countries, host.Country)
	}
	countries = uniqStrSlice(countries)
	sort.Strings(countries)
	return countries
}

// NodeList gets the node city items
func (d *Request) NodeList() []string {
	var node []string
	for _, host := range d.Hosts {
		if host.Country == d.Country {
			node = append(node, host.City)
		}
	}
	sort.Strings(node)
	return node
}

// ChkCountry validates and set requested country
func (d *Request) ChkCountry(country string) bool {
	country = strings.ToLower(country)
	for _, h := range d.Hosts {
		if country == h.Country {
			d.Country = country
			return true
		}
	}

	return false
}

// ChkNode set requested country
func (d *Request) ChkNode(city string) bool {
	city = strings.ToLower(city)
	for _, h := range d.Hosts {
		if d.Country == h.Country && city == h.City {
			d.Host = h.IP
			d.City = h.City
			return true
		}
	}
	return false
}

// Local set host to nothing means local
func (d *Request) Local() {
	d.Host = ""
	d.Country = ""
}

// Dig looks up name server w/ trace feature
func (d *Request) Dig() {
	if !d.TraceEnabled {
		d.RunDig()
	} else {
		d.RunDigTrace()
	}
}

// RunDig looks up name server
func (d *Request) RunDig() {
	c := new(dns.Client)
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(d.Target), d.Type)
	m.RecursionDesired = true
	m.RecursionAvailable = true

	println("Trying to query server:", d.Host, d.Country, d.City)

	r, rtt, err := c.Exchange(m, net.JoinHostPort(d.Host, "53"))
	if err != nil {
		println(err.Error())
		return
	}
	// Answer
	println(r.MsgHdr.String())
	for _, a := range r.Answer {
		fmt.Println(a)
	}
	// Extra info
	if len(r.Extra) > 0 {
		println("\n;; ADDITIONAL SECTION:")
		for _, a := range r.Extra {
			fmt.Println(a)
		}
	}
	fmt.Printf(";; Query time: %d ms\n", rtt/1e6)

	// CHAOS
	c.Timeout = ((rtt / 1e6) + 100) * time.Millisecond
	fmt.Printf("\n;; CHAOS CLASS BIND\n")
	for _, q := range []string{"version.bind.", "hostname.bind."} {
		m.Question[0] = dns.Question{q, dns.TypeTXT, dns.ClassCHAOS}
		r, _, err = c.Exchange(m, d.Host+":53")
		if err != nil {
			continue
		}
		for _, a := range r.Answer {
			fmt.Println(a)
		}
	}
}

// RunDigTrace handles dig trace
func (d *Request) RunDigTrace() {
	var (
		nss  = []string{d.Host}
		err  error
		host string
		rtt  time.Duration
		r    *dns.Msg
	)
	c := new(dns.Client)
	m := new(dns.Msg)
	m.RecursionDesired = true
	q := ""

	domain := []string{""}
	domain = append(domain, strings.Split(dns.Fqdn(d.Target), ".")...)
	for i := range domain {
		if i != 1 && i != len(domain)-1 {
			q = domain[len(domain)-i-1] + "." + q
		} else {
			q = domain[len(domain)-i-1] + q
		}

		if i != len(domain)-1 {
			m.SetQuestion(q, dns.TypeNS)
		} else {
			m.SetQuestion(q, d.Type)
		}

		for _, host = range nss {
			r, rtt, err = c.Exchange(m, net.JoinHostPort(host, "53"))
			if err != nil {
				println(err.Error())
			} else {
				break
			}
		}

		nss = nss[:0]

		for _, a := range r.Answer {
			println(a.String())
			if a.Header().Rrtype == dns.TypeNS {
				nss = append(nss, strings.Fields(a.String())[4])
			}
		}
		for _, a := range r.Ns {
			println(a.String())
			nss = append(nss, strings.Fields(a.String())[4])
		}

		fmt.Printf("from: %s#53 in %d ms\n", host, rtt/1e6)
	}
}

// cache provides caching for name servers
func (d *Request) cache(r string) bool {
	switch r {
	case "read":
		b, err := ioutil.ReadFile("/tmp/mylg.ns")
		if err != nil {
			panic(err.Error())
		}
		d.Hosts = d.Hosts[:0]
		r := bytes.NewBuffer(b)
		s := bufio.NewScanner(r)
		for s.Scan() {
			csv := strings.Split(s.Text(), ";")
			if len(csv) != 4 {
				continue
			}
			d.Hosts = append(d.Hosts, Host{Alpha2: csv[0], Country: csv[1], City: csv[2], IP: csv[3]})
		}
	case "write":
		var data []string
		for _, h := range d.Hosts {
			data = append(data, fmt.Sprintf("%s;%s;%s;%s", h.Alpha2, h.Country, h.City, h.IP))
		}
		err := ioutil.WriteFile("/tmp/mylg.ns", []byte(strings.Join(data, "\n")), 0644)
		if err != nil {
			panic(err.Error())
		}
	case "validate":
		f, err := os.Stat("/tmp/mylg.ns")
		if err != nil {
			return false
		}
		d := time.Since(f.ModTime())
		if d.Hours() > 48 {
			return false
		}
	}
	return true
}

// Fetch name servers from public-dns.info
func fetchNSHosts() []Host {
	var (
		hosts   []Host
		city    string
		counter = make(map[string]int)
		chkDup  = make(map[string]int)
	)
	resp, err := http.Get(publicDNSHost + publicDNSNodesPath)
	if err != nil {
		println(err.Error())
		return []Host{}
	}
	if resp.StatusCode != 200 {
		println("error: public dns is not available")
		return []Host{}
	}
	defer resp.Body.Close()
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		csv := strings.Split(scanner.Text(), ",")
		if csv[3] != "\"\"" {
			if name, ok := data.Country[csv[2]]; ok && counter[csv[2]] < 5 {
				name = strings.ToLower(name)
				city = strings.ToLower(csv[3])
				chkDup[name+city] += 1
				if chkDup[name+city] > 1 {
					city = fmt.Sprintf("%s0%d", city, chkDup[name+city]-1)
				}

				hosts = append(hosts, Host{IP: csv[0], Alpha2: csv[2], Country: name, City: city})
				counter[csv[2]]++
			}
		}
	}
	return hosts
}

// uniqStrSlice return unique slice
func uniqStrSlice(src []string) []string {
	var rst []string
	tmp := make(map[string]struct{})
	for _, s := range src {
		tmp[s] = struct{}{}
	}
	for s := range tmp {
		rst = append(rst, s)
	}
	return rst
}

// help
func help() {
	fmt.Println(`
    usage:
          dig [@local-server] host [options]
    options:
          +trace
    Example:
          dig google.com
          dig @8.8.8.8 yahoo.com
          dig google.com +trace
          dig google.com MX
	`)

}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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.
*/

// apiserver is the main api server and master for the cluster.
// it is responsible for serving the cluster management API.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os/user"
	"strings"
	"time"

	"code.google.com/p/goprotobuf/proto"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/record"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/master/ports"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/version/verflag"
	plugin "github.com/GoogleCloudPlatform/kubernetes/plugin/pkg/scheduler"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/mesos"
	_ "github.com/mesosphere/kubernetes-mesos/profile"
	kmscheduler "github.com/mesosphere/kubernetes-mesos/scheduler"
)

var (
	port                = flag.Int("port", ports.SchedulerPort, "The port that the scheduler's http service runs on")
	address             = util.IP(net.ParseIP("127.0.0.1"))
	etcdServerList      util.StringList
	etcdConfigFile      = flag.String("etcd_config", "", "The config file for the etcd client. Mutually exclusive with -etcd_servers.")
	clientConfig        = &client.Config{}
	mesosMaster         = flag.String("mesos_master", "localhost:5050", "Location of leading Mesos master. Default localhost:5050.")
	executorPath        = flag.String("executor_path", "", "Location of the kubernetes executor executable")
	proxyPath           = flag.String("proxy_path", "", "Location of the kubernetes proxy executable")
	mesosUser           = flag.String("mesos_user", "", "Mesos user for this framework, defaults to the username that owns the framework process.")
	mesosRole           = flag.String("mesos_role", "", "Mesos role for this framework, defaults to none.")
	mesosAuthPrincipal  = flag.String("mesos_authentication_principal", "", "Mesos authentication principal.")
	mesosAuthSecretFile = flag.String("mesos_authentication_secret_file", "", "Mesos authentication secret file.")
)

const (
	artifactPort     = 9000             // port of the service that services mesos artifacts (executor); TODO(jdef): make this configurable
	httpReadTimeout  = 10 * time.Second // k8s api server config: maximum duration before timing out read of the request
	httpWriteTimeout = 10 * time.Second // k8s api server config: maximum duration before timing out write of the response
)

func init() {
	flag.Var(&address, "address", "The IP address on to serve on (set to 0.0.0.0 for all interfaces). Default 127.0.0.1.")
	client.BindClientConfigFlags(flag.CommandLine, clientConfig)
}

// returns (downloadURI, basename(path))
func serveExecutorArtifact(path string) (*string, string) {
	serveFile := func(pattern string, filename string) {
		http.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, filename)
		})
	}

	// Create base path (http://foobar:5000/<base>)
	pathSplit := strings.Split(path, "/")
	var base string
	if len(pathSplit) > 0 {
		base = pathSplit[len(pathSplit)-1]
	} else {
		base = path
	}
	serveFile("/"+base, path)

	hostURI := fmt.Sprintf("http://%s:%d/%s", address.String(), artifactPort, base)
	log.V(2).Infof("Hosting artifact '%s' at '%s'", path, hostURI)

	return &hostURI, base
}

func prepareExecutorInfo() *mesos.ExecutorInfo {
	executorUris := []*mesos.CommandInfo_URI{}
	uri, _ := serveExecutorArtifact(*proxyPath)
	executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri, Executable: proto.Bool(true)})
	uri, executorCmd := serveExecutorArtifact(*executorPath)
	executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri, Executable: proto.Bool(true)})

	//TODO(jdef): provide some way (env var?) for user's to customize executor config
	//TODO(jdef): set -hostname_override and -address to 127.0.0.1 if `address` is 127.0.0.1
	//TODO(jdef): kubelet can publish events to the api server, we should probably tell it our IP address
	executorCommand := fmt.Sprintf("./%s -v=2 -hostname_override=0.0.0.0 -allow_privileged=%t", executorCmd, *allowPrivileged)
	if len(etcdServerList) > 0 {
		etcdServerArguments := strings.Join(etcdServerList, ",")
		executorCommand = fmt.Sprintf("%s -etcd_servers=%s", executorCommand, etcdServerArguments)
	} else {
		uri, basename := serveExecutorArtifact(*etcdConfigFile)
		executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri})
		executorCommand = fmt.Sprintf("%s -etcd_config=./%s", executorCommand, basename)
	}

	go http.ListenAndServe(fmt.Sprintf("%s:%d", address.String(), artifactPort), nil)
	log.V(2).Info("Serving executor artifacts...")

	// Create mesos scheduler driver.
	return &mesos.ExecutorInfo{
		ExecutorId: &mesos.ExecutorID{Value: proto.String("KubeleteExecutorID")},
		Command: &mesos.CommandInfo{
			Value: proto.String(executorCommand),
			Uris:  executorUris,
		},
		Name:   proto.String("Kubelet Executor"),
		Source: proto.String("kubernetes"),
	}
}

// Copied from cmd/apiserver.go
func main() {
	flag.Parse()
	util.InitLogs()
	defer util.FlushLogs()

	verflag.PrintAndExitIfRequested()

	// we'll need this for the kubelet-executor
	if (*etcdConfigFile != "" && len(etcdServerList) != 0) || (*etcdConfigFile == "" && len(etcdServerList) == 0) {
		log.Fatalf("specify either -etcd_servers or -etcd_config")
	}

	client, err := client.New(clientConfig)
	if err != nil {
		log.Fatalf("Invalid API configuration: %v", err)
	}

	record.StartRecording(client.Events(""), "scheduler")

	// Create mesos scheduler driver.
	executor := prepareExecutorInfo()
	mesosPodScheduler := kmscheduler.New(executor, kmscheduler.FCFSScheduleFunc, client)
	info, cred, err := buildFrameworkInfo()
	if err != nil {
		log.Fatalf("Misconfigured mesos framework: %v", err)
	}
	driver := &mesos.MesosSchedulerDriver{
		Master:    *mesosMaster,
		Framework: *info,
		Scheduler: mesosPodScheduler,
		Cred:      cred,
	}

	mesosPodScheduler.Init(driver)
	driver.Init()
	defer driver.Destroy()

	go func() {
		if st, err := driver.Start(); err == nil {
			if st != mesos.Status_DRIVER_RUNNING {
				log.Fatalf("Scheduler driver failed to start, has status: %v", st)
			}
			if st, err = driver.Join(); err != nil {
				log.Fatal(err)
			} else if st != mesos.Status_DRIVER_RUNNING {
				log.Fatalf("Scheduler driver failed to join, has status: %v", st)
			}
		} else {
			log.Fatalf("Failed to start driver: %v", err)
		}
	}()

	plugin.New(mesosPodScheduler.NewPluginConfig()).Run()

	select {}
}

func buildFrameworkInfo() (info *mesos.FrameworkInfo, cred *mesos.Credential, err error) {

	username, err := getUsername()
	if err != nil {
		return nil, nil, err
	}
	log.V(2).Infof("Framework configured with mesos user %v", username)
	info = &mesos.FrameworkInfo{
		Name: proto.String("KubernetesScheduler"),
		User: proto.String(username),
	}
	if *mesosRole != "" {
		info.Role = proto.String(*mesosRole)
	}
	if *mesosAuthPrincipal != "" {
		info.Principal = proto.String(*mesosAuthPrincipal)
		secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
		if err != nil {
			return nil, nil, err
		}
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
			Secret:    secret,
		}
	}
	return
}

func getUsername() (username string, err error) {
	username = *mesosUser
	if username == "" {
		if u, err := user.Current(); err == nil {
			username = u.Username
			if username == "" {
				username = "root"
			}
		}
	}
	return
}
<commit_msg>segregated code for readability, in preparation for eliminating pod registry implementation<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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.
*/

// apiserver is the main api server and master for the cluster.
// it is responsible for serving the cluster management API.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os/user"
	"strconv"
	"strings"
	"time"

	"code.google.com/p/goprotobuf/proto"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/record"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/master/ports"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/version/verflag"
	plugin "github.com/GoogleCloudPlatform/kubernetes/plugin/pkg/scheduler"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/mesos"
	_ "github.com/mesosphere/kubernetes-mesos/profile"
	kmscheduler "github.com/mesosphere/kubernetes-mesos/scheduler"
)

var (
	port                = flag.Int("port", ports.SchedulerPort, "The port that the scheduler's http service runs on")
	address             = util.IP(net.ParseIP("127.0.0.1"))
	etcdServerList      util.StringList
	etcdConfigFile      = flag.String("etcd_config", "", "The config file for the etcd client. Mutually exclusive with -etcd_servers.")
	clientConfig        = &client.Config{}
	allowPrivileged     = flag.Bool("allow_privileged", false, "If true, allow privileged containers. Default false.")
	mesosMaster         = flag.String("mesos_master", "localhost:5050", "Location of leading Mesos master. Default localhost:5050.")
	executorPath        = flag.String("executor_path", "", "Location of the kubernetes executor executable")
	proxyPath           = flag.String("proxy_path", "", "Location of the kubernetes proxy executable")
	mesosUser           = flag.String("mesos_user", "", "Mesos user for this framework, defaults to the username that owns the framework process.")
	mesosRole           = flag.String("mesos_role", "", "Mesos role for this framework, defaults to none.")
	mesosAuthPrincipal  = flag.String("mesos_authentication_principal", "", "Mesos authentication principal.")
	mesosAuthSecretFile = flag.String("mesos_authentication_secret_file", "", "Mesos authentication secret file.")
)

const (
	artifactPort     = 9000             // port of the service that services mesos artifacts (executor); TODO(jdef): make this configurable
	httpReadTimeout  = 10 * time.Second // k8s api server config: maximum duration before timing out read of the request
	httpWriteTimeout = 10 * time.Second // k8s api server config: maximum duration before timing out write of the response
)

func init() {
	flag.Var(&address, "address", "The IP address on to serve on (set to 0.0.0.0 for all interfaces). Default 127.0.0.1.")
	client.BindClientConfigFlags(flag.CommandLine, clientConfig)
}

// returns (downloadURI, basename(path))
func serveExecutorArtifact(path string) (*string, string) {
	serveFile := func(pattern string, filename string) {
		http.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, filename)
		})
	}

	// Create base path (http://foobar:5000/<base>)
	pathSplit := strings.Split(path, "/")
	var base string
	if len(pathSplit) > 0 {
		base = pathSplit[len(pathSplit)-1]
	} else {
		base = path
	}
	serveFile("/"+base, path)

	hostURI := fmt.Sprintf("http://%s:%d/%s", address.String(), artifactPort, base)
	log.V(2).Infof("Hosting artifact '%s' at '%s'", path, hostURI)

	return &hostURI, base
}

func prepareExecutorInfo() *mesos.ExecutorInfo {
	executorUris := []*mesos.CommandInfo_URI{}
	uri, _ := serveExecutorArtifact(*proxyPath)
	executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri, Executable: proto.Bool(true)})
	uri, executorCmd := serveExecutorArtifact(*executorPath)
	executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri, Executable: proto.Bool(true)})

	//TODO(jdef): provide some way (env var?) for user's to customize executor config
	//TODO(jdef): set -hostname_override and -address to 127.0.0.1 if `address` is 127.0.0.1
	//TODO(jdef): kubelet can publish events to the api server, we should probably tell it our IP address
	executorCommand := fmt.Sprintf("./%s -v=2 -hostname_override=0.0.0.0 -allow_privileged=%t", executorCmd, *allowPrivileged)
	if len(etcdServerList) > 0 {
		etcdServerArguments := strings.Join(etcdServerList, ",")
		executorCommand = fmt.Sprintf("%s -etcd_servers=%s", executorCommand, etcdServerArguments)
	} else {
		uri, basename := serveExecutorArtifact(*etcdConfigFile)
		executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri})
		executorCommand = fmt.Sprintf("%s -etcd_config=./%s", executorCommand, basename)
	}

	go http.ListenAndServe(fmt.Sprintf("%s:%d", address.String(), artifactPort), nil)
	log.V(2).Info("Serving executor artifacts...")

	// Create mesos scheduler driver.
	return &mesos.ExecutorInfo{
		ExecutorId: &mesos.ExecutorID{Value: proto.String("KubeleteExecutorID")},
		Command: &mesos.CommandInfo{
			Value: proto.String(executorCommand),
			Uris:  executorUris,
		},
		Name:   proto.String("Kubelet Executor"),
		Source: proto.String("kubernetes"),
	}
}

// Copied from cmd/apiserver.go
func main() {
	flag.Parse()
	util.InitLogs()
	defer util.FlushLogs()

	verflag.PrintAndExitIfRequested()

	// we'll need this for the kubelet-executor
	if (*etcdConfigFile != "" && len(etcdServerList) != 0) || (*etcdConfigFile == "" && len(etcdServerList) == 0) {
		log.Fatalf("specify either -etcd_servers or -etcd_config")
	}

	client, err := client.New(clientConfig)
	if err != nil {
		log.Fatalf("Invalid API configuration: %v", err)
	}

	record.StartRecording(client.Events(""), "scheduler")

	// Create mesos scheduler driver.
	executor := prepareExecutorInfo()
	mesosPodScheduler := kmscheduler.New(executor, kmscheduler.FCFSScheduleFunc, client)
	info, cred, err := buildFrameworkInfo()
	if err != nil {
		log.Fatalf("Misconfigured mesos framework: %v", err)
	}
	driver := &mesos.MesosSchedulerDriver{
		Master:    *mesosMaster,
		Framework: *info,
		Scheduler: mesosPodScheduler,
		Cred:      cred,
	}

	mesosPodScheduler.Init(driver)
	driver.Init()
	defer driver.Destroy()

	go func() {
		if st, err := driver.Start(); err == nil {
			if st != mesos.Status_DRIVER_RUNNING {
				log.Fatalf("Scheduler driver failed to start, has status: %v", st)
			}
			if st, err = driver.Join(); err != nil {
				log.Fatal(err)
			} else if st != mesos.Status_DRIVER_RUNNING {
				log.Fatalf("Scheduler driver failed to join, has status: %v", st)
			}
		} else {
			log.Fatalf("Failed to start driver: %v", err)
		}
	}()

	go http.ListenAndServe(net.JoinHostPort(address.String(), strconv.Itoa(*port)), nil)
	plugin.New(mesosPodScheduler.NewPluginConfig()).Run()

	select {}
}

func buildFrameworkInfo() (info *mesos.FrameworkInfo, cred *mesos.Credential, err error) {

	username, err := getUsername()
	if err != nil {
		return nil, nil, err
	}
	log.V(2).Infof("Framework configured with mesos user %v", username)
	info = &mesos.FrameworkInfo{
		Name: proto.String("KubernetesScheduler"),
		User: proto.String(username),
	}
	if *mesosRole != "" {
		info.Role = proto.String(*mesosRole)
	}
	if *mesosAuthPrincipal != "" {
		info.Principal = proto.String(*mesosAuthPrincipal)
		secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
		if err != nil {
			return nil, nil, err
		}
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
			Secret:    secret,
		}
	}
	return
}

func getUsername() (username string, err error) {
	username = *mesosUser
	if username == "" {
		if u, err := user.Current(); err == nil {
			username = u.Username
			if username == "" {
				username = "root"
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2010 Jeremy Wall (jeremy@marzhillstudios.com)
// Use of this source code is governed by the Artistic License 2.0.
// That License is included in the LICENSE file.

package transform

import (
	"code.google.com/p/go-html-transform/css/selector"
	"code.google.com/p/go-html-transform/h5"

	"code.google.com/p/go.net/html"
	"log"
)

// Collector defines an interface for html node collectors.
type Collector interface {
	// Find searches a tree rooted at n and returns a slice of nodes
	// that match a criteria.
	Find(n *html.Node) []*html.Node
}

// The TransformFunc type is the type of a html.Node transformation function.
type TransformFunc func(*html.Node)

// Transformer encapsulates a document under transformation.
type Transformer struct {
	doc h5.Tree
}

// Constructor for a Transformer. It makes a copy of the document
// and transforms that instead of the original.
func NewTransformer(t h5.Tree) *Transformer {
	return newTransformer(t.Clone())
}

func newTransformer(t h5.Tree) *Transformer {
	return &Transformer{doc: t}
}

// The Doc method returns the document under transformation.
func (t *Transformer) Doc() *html.Node {
	return t.doc.Top()
}

func (t *Transformer) String() string {
	return t.doc.String()
}

func (t *Transformer) Clone() *Transformer {
	return NewTransformer(t.doc)
}

func applyFuncToCollector(f TransformFunc, n *html.Node, sel Collector) {
	for _, nn := range sel.Find(n) {
		f(nn)
	}
}

// The ApplyWithSelector method applies a TransformFunc to the nodes matched
// by the CSS3 Selector.
func (t *Transformer) Apply(f TransformFunc, sel string) error {
	sq, err := selector.Selector(sel)
	t.ApplyWithCollector(f, sq)
	return err
}

// ApplyWithCollector applies a TransformFunc to the tree using a Collector.
func (t *Transformer) ApplyWithCollector(f TransformFunc, coll Collector) {
	// TODO come up with a way to walk tree once?
	applyFuncToCollector(f, t.Doc(), coll)
}

// Transform is a bundle of selectors and a transform func. It forms a
// self contained Transfrom on an html document that can be reused.
type Transform struct {
	coll Collector
	f    TransformFunc
}

// Trans creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a valid CSS3 Selector.
// It returns a *Transform or an error if the selector wasn't valid
func Trans(f TransformFunc, sel string) (*Transform, error) {
	sq, err := selector.Selector(sel)
	return TransCollector(f, sq), err
}

// MustTrans creates a Transform.
// Panics if the selector wasn't valid.
func MustTrans(f TransformFunc, sel string) *Transform {
	t, err := Trans(f, sel)
	if err != nil { panic(err) }
	return t
}

// TransCollector creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a Collector
func TransCollector(f TransformFunc, coll Collector) *Transform {
	return &Transform{f: f, coll: coll}
}

// ApplyAll applies a series of Transforms to a document.
//     t.ApplyAll(Trans(f, sel1, sel2), Trans(f2, sel3, sel4))
func (t *Transformer) ApplyAll(ts ...*Transform) {
	for _, spec := range ts {
		t.ApplyWithCollector(spec.f, spec.coll)
	}
}

// AppendChildren creates a TransformFunc that appends the Children passed in.
func AppendChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			if c.Parent != nil {
				c.Parent.RemoveChild(c)
			}
			n.AppendChild(c)
		}
	}
}

// PrependChildren creates a TransformFunc that prepends the Children passed in.
func PrependChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			n.InsertBefore(c, n.FirstChild)
		}
	}
}

// RemoveChildren creates a TransformFunc that removes the Children of the node
// it operates on.
func RemoveChildren() TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
	}
}

func removeChildren(n *html.Node) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		defer n.RemoveChild(c)
	}
}

// ReplaceChildren creates a TransformFunc that replaces the Children of the
// node it operates on with the Children passed in.
func ReplaceChildren(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
		for _, c := range ns {
			n.AppendChild(c)
		}
	}
}

func nodeToString(n *html.Node) string {
	t := h5.NewTree(n)
	return t.String()
}

// Replace constructs a TransformFunc that replaces a node with the nodes passed
// in.
func Replace(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		p := n.Parent
		switch p {
		case nil:
			log.Panicf("Attempt to replace Root node: %s", h5.RenderNodesToString([]*html.Node{n}))
		default:
			for _, nc := range ns {
				p.InsertBefore(nc, n)
			}
			p.RemoveChild(n)
		}
	}
}

// DoAll returns a TransformFunc that combines all the TransformFuncs that are
// passed in. Doing each transform in order.
func DoAll(fs ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, f := range fs {
			f(n)
		}
	}
}

// CopyAnd will construct a TransformFunc that will
// make a copy of the node for each passed in TransformFunc
// and replace the passed in node with the resulting transformed
// html.Nodes.
func CopyAnd(fns ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, fn := range fns {
			node := h5.CloneNode(n)
			n.Parent.InsertBefore(node, n)
			fn(node)
		}
		n.Parent.RemoveChild(n)
	}
}

// SubTransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func Subtransform(f TransformFunc, sel string) (TransformFunc, error) {
	sq, err := selector.Selector(sel)
	return SubtransformCollector(f, sq), err
}

// MustSubtransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// Panics if the selector string is malformed.
func MustSubtransform(f TransformFunc, sel string) TransformFunc {
	t, err := Subtransform(f, sel)
	if err != nil { panic(err) }
	return t
}

// SubTransformSelector constructs a TransformFunc that runs a TransformFunc on
// any nodes collected, using the passed in collector, from the subtree the
// TransformFunc is run on.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func SubtransformCollector(f TransformFunc, coll Collector) TransformFunc {
	return func(n *html.Node) {
		applyFuncToCollector(f, n, coll)
	}
}

// ModifyAttrb creates a TransformFunc that modifies the attributes
// of the node it operates on. If an Attribute with the same name
// as the key doesn't exist it creates it.
func ModifyAttrib(key string, val string) TransformFunc {
	return func(n *html.Node) {
		found := false
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = val
				found = true
			}
		}
		if !found {
			n.Attr = append(n.Attr, html.Attribute{Key: key, Val: val})
		}
	}
}

// TransformAttrib returns a TransformFunc that transforms an attribute on
// the node it operates on using the provided func. It only transforms
// the attribute if it exists.
func TransformAttrib(key string, f func(string) string) TransformFunc {
	return func(n *html.Node) {
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = f(n.Attr[i].Val)
			}
		}
	}
}

// Trace is a debugging wrapper for transform funcs.
// It prints debugging information before and after the TransformFunc
// is applied.
func Trace(f TransformFunc, msg string, args ...interface{}) TransformFunc {
	return func(n *html.Node) {
		log.Printf("TRACE: "+msg, args...)
		p := n.Parent
		if p == nil {
			p = n
		}
		log.Printf("TRACE: Before: %s", h5.NewTree(p).String())
		f(n)
		log.Printf("TRACE: After: %s", h5.NewTree(p).String())
	}
}
<commit_msg>Add Render method for the html transform documents.<commit_after>// Copyright 2010 Jeremy Wall (jeremy@marzhillstudios.com)
// Use of this source code is governed by the Artistic License 2.0.
// That License is included in the LICENSE file.

package transform

import (
	"io"
	"log"

	"code.google.com/p/go.net/html"

	"code.google.com/p/go-html-transform/css/selector"
	"code.google.com/p/go-html-transform/h5"
)

// Collector defines an interface for html node collectors.
type Collector interface {
	// Find searches a tree rooted at n and returns a slice of nodes
	// that match a criteria.
	Find(n *html.Node) []*html.Node
}

// The TransformFunc type is the type of a html.Node transformation function.
type TransformFunc func(*html.Node)

// Transformer encapsulates a document under transformation.
type Transformer struct {
	doc h5.Tree
}

// Constructor for a Transformer. It makes a copy of the document
// and transforms that instead of the original.
func NewTransformer(t h5.Tree) *Transformer {
	return newTransformer(t.Clone())
}

func newTransformer(t h5.Tree) *Transformer {
	return &Transformer{doc: t}
}

// The Doc method returns the document under transformation.
func (t *Transformer) Doc() *html.Node {
	return t.doc.Top()
}

func (t *Transformer) Render(w io.Writer) error {
	return t.doc.Render(w)
}

func (t *Transformer) String() string {
	return t.doc.String()
}

func (t *Transformer) Clone() *Transformer {
	return NewTransformer(t.doc)
}

func applyFuncToCollector(f TransformFunc, n *html.Node, sel Collector) {
	for _, nn := range sel.Find(n) {
		f(nn)
	}
}

// The ApplyWithSelector method applies a TransformFunc to the nodes matched
// by the CSS3 Selector.
func (t *Transformer) Apply(f TransformFunc, sel string) error {
	sq, err := selector.Selector(sel)
	t.ApplyWithCollector(f, sq)
	return err
}

// ApplyWithCollector applies a TransformFunc to the tree using a Collector.
func (t *Transformer) ApplyWithCollector(f TransformFunc, coll Collector) {
	// TODO come up with a way to walk tree once?
	applyFuncToCollector(f, t.Doc(), coll)
}

// Transform is a bundle of selectors and a transform func. It forms a
// self contained Transfrom on an html document that can be reused.
type Transform struct {
	coll Collector
	f    TransformFunc
}

// Trans creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a valid CSS3 Selector.
// It returns a *Transform or an error if the selector wasn't valid
func Trans(f TransformFunc, sel string) (*Transform, error) {
	sq, err := selector.Selector(sel)
	return TransCollector(f, sq), err
}

// MustTrans creates a Transform.
// Panics if the selector wasn't valid.
func MustTrans(f TransformFunc, sel string) *Transform {
	t, err := Trans(f, sel)
	if err != nil {
		panic(err)
	}
	return t
}

// TransCollector creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a Collector
func TransCollector(f TransformFunc, coll Collector) *Transform {
	return &Transform{f: f, coll: coll}
}

// ApplyAll applies a series of Transforms to a document.
//     t.ApplyAll(Trans(f, sel1, sel2), Trans(f2, sel3, sel4))
func (t *Transformer) ApplyAll(ts ...*Transform) {
	for _, spec := range ts {
		t.ApplyWithCollector(spec.f, spec.coll)
	}
}

// AppendChildren creates a TransformFunc that appends the Children passed in.
func AppendChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			if c.Parent != nil {
				c.Parent.RemoveChild(c)
			}
			n.AppendChild(c)
		}
	}
}

// PrependChildren creates a TransformFunc that prepends the Children passed in.
func PrependChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			n.InsertBefore(c, n.FirstChild)
		}
	}
}

// RemoveChildren creates a TransformFunc that removes the Children of the node
// it operates on.
func RemoveChildren() TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
	}
}

func removeChildren(n *html.Node) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		defer n.RemoveChild(c)
	}
}

// ReplaceChildren creates a TransformFunc that replaces the Children of the
// node it operates on with the Children passed in.
func ReplaceChildren(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
		for _, c := range ns {
			n.AppendChild(c)
		}
	}
}

func nodeToString(n *html.Node) string {
	t := h5.NewTree(n)
	return t.String()
}

// Replace constructs a TransformFunc that replaces a node with the nodes passed
// in.
func Replace(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		p := n.Parent
		switch p {
		case nil:
			log.Panicf("Attempt to replace Root node: %s", h5.RenderNodesToString([]*html.Node{n}))
		default:
			for _, nc := range ns {
				p.InsertBefore(nc, n)
			}
			p.RemoveChild(n)
		}
	}
}

// DoAll returns a TransformFunc that combines all the TransformFuncs that are
// passed in. Doing each transform in order.
func DoAll(fs ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, f := range fs {
			f(n)
		}
	}
}

// CopyAnd will construct a TransformFunc that will
// make a copy of the node for each passed in TransformFunc
// and replace the passed in node with the resulting transformed
// html.Nodes.
func CopyAnd(fns ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, fn := range fns {
			node := h5.CloneNode(n)
			n.Parent.InsertBefore(node, n)
			fn(node)
		}
		n.Parent.RemoveChild(n)
	}
}

// SubTransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func Subtransform(f TransformFunc, sel string) (TransformFunc, error) {
	sq, err := selector.Selector(sel)
	return SubtransformCollector(f, sq), err
}

// MustSubtransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// Panics if the selector string is malformed.
func MustSubtransform(f TransformFunc, sel string) TransformFunc {
	t, err := Subtransform(f, sel)
	if err != nil {
		panic(err)
	}
	return t
}

// SubTransformSelector constructs a TransformFunc that runs a TransformFunc on
// any nodes collected, using the passed in collector, from the subtree the
// TransformFunc is run on.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func SubtransformCollector(f TransformFunc, coll Collector) TransformFunc {
	return func(n *html.Node) {
		applyFuncToCollector(f, n, coll)
	}
}

// ModifyAttrb creates a TransformFunc that modifies the attributes
// of the node it operates on. If an Attribute with the same name
// as the key doesn't exist it creates it.
func ModifyAttrib(key string, val string) TransformFunc {
	return func(n *html.Node) {
		found := false
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = val
				found = true
			}
		}
		if !found {
			n.Attr = append(n.Attr, html.Attribute{Key: key, Val: val})
		}
	}
}

// TransformAttrib returns a TransformFunc that transforms an attribute on
// the node it operates on using the provided func. It only transforms
// the attribute if it exists.
func TransformAttrib(key string, f func(string) string) TransformFunc {
	return func(n *html.Node) {
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = f(n.Attr[i].Val)
			}
		}
	}
}

// Trace is a debugging wrapper for transform funcs.
// It prints debugging information before and after the TransformFunc
// is applied.
func Trace(f TransformFunc, msg string, args ...interface{}) TransformFunc {
	return func(n *html.Node) {
		log.Printf("TRACE: "+msg, args...)
		p := n.Parent
		if p == nil {
			p = n
		}
		log.Printf("TRACE: Before: %s", h5.NewTree(p).String())
		f(n)
		log.Printf("TRACE: After: %s", h5.NewTree(p).String())
	}
}
<|endoftext|>


<commit_before>package api

import (
	"time"
)

// ImageExportPost represents the fields required to export a LXD image
//
// swagger:model
//
// API extension: images_push_relay
type ImageExportPost struct {
	// Target server URL
	// Example: https://1.2.3.4:8443
	Target string `json:"target" yaml:"target"`

	// Image receive secret
	// Example: RANDOM-STRING
	Secret string `json:"secret" yaml:"secret"`

	// Remote server certificate
	// Example: X509 PEM certificate
	Certificate string `json:"certificate" yaml:"certificate"`

	// List of aliases to set on the image
	Aliases []ImageAlias `json:"aliases" yaml:"aliases"`
}

// ImagesPost represents the fields available for a new LXD image
//
// swagger:model
type ImagesPost struct {
	ImagePut `yaml:",inline"`

	// Original filename of the image
	// Example: lxd.tar.xz
	Filename string `json:"filename" yaml:"filename"`

	// Source of the image
	Source *ImagesPostSource `json:"source" yaml:"source"`

	// Compression algorithm to use when turning an instance into an image
	// Example: gzip
	//
	// API extension: image_compression_algorithm
	CompressionAlgorithm string `json:"compression_algorithm" yaml:"compression_algorithm"`

	// Aliases to add to the image
	// Example: [{"name": "foo"}, {"name": "bar"}]
	//
	// API extension: image_create_aliases
	Aliases []ImageAlias `json:"aliases" yaml:"aliases"`
}

// ImagesPostSource represents the source of a new LXD image
//
// swagger:model
type ImagesPostSource struct {
	ImageSource `yaml:",inline"`

	// Transfer mode (push or pull)
	// Example: pull
	Mode string `json:"mode" yaml:"mode"`

	// Type of image source (instance, snapshot, image or url)
	// Example: instance
	Type string `json:"type" yaml:"type"`

	// Source URL (for type "url")
	// Example: https://some-server.com/some-directory/
	URL string `json:"url" yaml:"url"`

	// Instance name (for type "instance" or "snapshot")
	// Example: c1/snap0
	Name string `json:"name" yaml:"name"`

	// Source image fingerprint (for type "image")
	// Example: 8ae945c52bb2f2df51c923b04022312f99bbb72c356251f54fa89ea7cf1df1d0
	Fingerprint string `json:"fingerprint" yaml:"fingerprint"`

	// Source image server secret token (when downloading private images)
	// Example: RANDOM-STRING
	Secret string `json:"secret" yaml:"secret"`

	// Source project name
	// Example: project1
	//
	// API extension: image_source_project
	Project string `json:"project" yaml:"project"`
}

// ImagePut represents the modifiable fields of a LXD image
//
// swagger:model
type ImagePut struct {
	// Whether the image should auto-update when a new build is available
	// Example: true
	AutoUpdate bool `json:"auto_update" yaml:"auto_update"`

	// Descriptive properties
	// Example: {"os": "Ubuntu", "release": "focal", "variant": "cloud"}
	Properties map[string]string `json:"properties" yaml:"properties"`

	// Whether the image is available to unauthenticated users
	// Example: false
	Public bool `json:"public" yaml:"public"`

	// When the image becomes obsolete
	// Example: 2025-03-23T20:00:00-04:00
	//
	// API extension: images_expiry
	ExpiresAt time.Time `json:"expires_at" yaml:"expires_at"`

	// List of profiles to use when creating from this image (if none provided by user)
	// Example: ["default"]
	//
	// API extension: image_profiles
	Profiles []string `json:"profiles" yaml:"profiles"`
}

// Image represents a LXD image
//
// swagger:model
type Image struct {
	ImagePut `yaml:",inline"`

	// List of aliases
	Aliases []ImageAlias `json:"aliases" yaml:"aliases"`

	// Architecture
	// Example: x86_64
	Architecture string `json:"architecture" yaml:"architecture"`

	// Whether the image is an automatically cached remote image
	// Example: true
	Cached bool `json:"cached" yaml:"cached"`

	// Original filename
	// Example: 06b86454720d36b20f94e31c6812e05ec51c1b568cf3a8abd273769d213394bb.rootfs
	Filename string `json:"filename" yaml:"filename"`

	// Full SHA-256 fingerprint
	// Example: 06b86454720d36b20f94e31c6812e05ec51c1b568cf3a8abd273769d213394bb
	Fingerprint string `json:"fingerprint" yaml:"fingerprint"`

	// Size of the image in bytes
	// Example: 272237676
	Size int64 `json:"size" yaml:"size"`

	// Where the image came from
	UpdateSource *ImageSource `json:"update_source,omitempty" yaml:"update_source,omitempty"`

	// Type of image (container or virtual-machine)
	// Example: container
	//
	// API extension: image_types
	Type string `json:"type" yaml:"type"`

	// When the image was originally created
	// Example: 2021-03-23T20:00:00-04:00
	CreatedAt time.Time `json:"created_at" yaml:"created_at"`

	// Last time the image was used
	// Example: 2021-03-22T20:39:00.575185384-04:00
	LastUsedAt time.Time `json:"last_used_at" yaml:"last_used_at"`

	// When the image was added to this LXD server
	// Example: 2021-03-24T14:18:15.115036787-04:00
	UploadedAt time.Time `json:"uploaded_at" yaml:"uploaded_at"`
}

// Writable converts a full Image struct into a ImagePut struct (filters read-only fields)
func (img *Image) Writable() ImagePut {
	return img.ImagePut
}

// ImageAlias represents an alias from the alias list of a LXD image
//
// swagger:model
type ImageAlias struct {
	// Name of the alias
	// Example: ubuntu-20.04
	Name string `json:"name" yaml:"name"`

	// Description of the alias
	// Example: Our preferred Ubuntu image
	Description string `json:"description" yaml:"description"`
}

// ImageSource represents the source of a LXD image
//
// swagger:model
type ImageSource struct {
	// Source alias to download from
	// Example: focal
	Alias string `json:"alias" yaml:"alias"`

	// Source server certificate (if not trusted by system CA)
	// Example: X509 PEM certificate
	Certificate string `json:"certificate" yaml:"certificate"`

	// Source server protocol
	// Example: simplestreams
	Protocol string `json:"protocol" yaml:"protocol"`

	// URL of the source server
	// Example: https://images.linuxcontainers.org
	Server string `json:"server" yaml:"server"`

	// Type of image (container or virtual-machine)
	// Example: container
	//
	// API extension: image_types
	ImageType string `json:"image_type" yaml:"image_type"`
}

// ImageAliasesPost represents a new LXD image alias
//
// swagger:model
type ImageAliasesPost struct {
	ImageAliasesEntry `yaml:",inline"`
}

// ImageAliasesEntryPost represents the required fields to rename a LXD image alias
//
// swagger:model
type ImageAliasesEntryPost struct {
	// Alias name
	// Example: ubuntu-20.04
	Name string `json:"name" yaml:"name"`
}

// ImageAliasesEntryPut represents the modifiable fields of a LXD image alias
//
// swagger:model
type ImageAliasesEntryPut struct {
	// Alias description
	// Example: Our preferred Ubuntu image
	Description string `json:"description" yaml:"description"`

	// Target fingerprint for the alias
	// Example: 06b86454720d36b20f94e31c6812e05ec51c1b568cf3a8abd273769d213394bb
	Target string `json:"target" yaml:"target"`
}

// ImageAliasesEntry represents a LXD image alias
//
// swagger:model
type ImageAliasesEntry struct {
	ImageAliasesEntryPut `yaml:",inline"`

	// Alias name
	// Example: ubuntu-20.04
	Name string `json:"name" yaml:"name"`

	// Alias type (container or virtual-machine)
	// Example: container
	//
	// API extension: image_types
	Type string `json:"type" yaml:"type"`
}

// ImageMetadata represents LXD image metadata (used in image tarball)
//
// swagger:model
type ImageMetadata struct {
	// Architecture name
	// Example: x86_64
	Architecture string `json:"architecture" yaml:"architecture"`

	// Image creation data (as UNIX epoch)
	// Example: 1620655439
	CreationDate int64 `json:"creation_date" yaml:"creation_date"`

	// Image expiry data (as UNIX epoch)
	// Example: 1620685757
	ExpiryDate int64 `json:"expiry_date" yaml:"expiry_date"`

	// Descriptive properties
	// Example: {"os": "Ubuntu", "release": "focal", "variant": "cloud"}
	Properties map[string]string `json:"properties" yaml:"properties"`

	// Template for files in the image
	Templates map[string]*ImageMetadataTemplate `json:"templates" yaml:"templates"`
}

// ImageMetadataTemplate represents a template entry in image metadata (used in image tarball)
//
// swagger:model
type ImageMetadataTemplate struct {
	// When to trigger the template (create, copy or start)
	// Example: create
	When []string `json:"when" yaml:"when"`

	// Whether to trigger only if the file is missing
	// Example: false
	CreateOnly bool `json:"create_only" yaml:"create_only"`

	// The template itself as a valid pongo2 template
	// Example: pongo2-template
	Template string `json:"template" yaml:"template"`

	// Key/value properties to pass to the template
	// Example: {"foo": "bar"}
	Properties map[string]string `json:"properties" yaml:"properties"`
}
<commit_msg>shared/api: Add Project to ImageExportPost<commit_after>package api

import (
	"time"
)

// ImageExportPost represents the fields required to export a LXD image
//
// swagger:model
//
// API extension: images_push_relay
type ImageExportPost struct {
	// Target server URL
	// Example: https://1.2.3.4:8443
	Target string `json:"target" yaml:"target"`

	// Image receive secret
	// Example: RANDOM-STRING
	Secret string `json:"secret" yaml:"secret"`

	// Remote server certificate
	// Example: X509 PEM certificate
	Certificate string `json:"certificate" yaml:"certificate"`

	// List of aliases to set on the image
	Aliases []ImageAlias `json:"aliases" yaml:"aliases"`

	// Project name
	// Example: project1
	//
	// API extension: image_target_project
	Project string `json:"project" yaml:"project"`
}

// ImagesPost represents the fields available for a new LXD image
//
// swagger:model
type ImagesPost struct {
	ImagePut `yaml:",inline"`

	// Original filename of the image
	// Example: lxd.tar.xz
	Filename string `json:"filename" yaml:"filename"`

	// Source of the image
	Source *ImagesPostSource `json:"source" yaml:"source"`

	// Compression algorithm to use when turning an instance into an image
	// Example: gzip
	//
	// API extension: image_compression_algorithm
	CompressionAlgorithm string `json:"compression_algorithm" yaml:"compression_algorithm"`

	// Aliases to add to the image
	// Example: [{"name": "foo"}, {"name": "bar"}]
	//
	// API extension: image_create_aliases
	Aliases []ImageAlias `json:"aliases" yaml:"aliases"`
}

// ImagesPostSource represents the source of a new LXD image
//
// swagger:model
type ImagesPostSource struct {
	ImageSource `yaml:",inline"`

	// Transfer mode (push or pull)
	// Example: pull
	Mode string `json:"mode" yaml:"mode"`

	// Type of image source (instance, snapshot, image or url)
	// Example: instance
	Type string `json:"type" yaml:"type"`

	// Source URL (for type "url")
	// Example: https://some-server.com/some-directory/
	URL string `json:"url" yaml:"url"`

	// Instance name (for type "instance" or "snapshot")
	// Example: c1/snap0
	Name string `json:"name" yaml:"name"`

	// Source image fingerprint (for type "image")
	// Example: 8ae945c52bb2f2df51c923b04022312f99bbb72c356251f54fa89ea7cf1df1d0
	Fingerprint string `json:"fingerprint" yaml:"fingerprint"`

	// Source image server secret token (when downloading private images)
	// Example: RANDOM-STRING
	Secret string `json:"secret" yaml:"secret"`

	// Source project name
	// Example: project1
	//
	// API extension: image_source_project
	Project string `json:"project" yaml:"project"`
}

// ImagePut represents the modifiable fields of a LXD image
//
// swagger:model
type ImagePut struct {
	// Whether the image should auto-update when a new build is available
	// Example: true
	AutoUpdate bool `json:"auto_update" yaml:"auto_update"`

	// Descriptive properties
	// Example: {"os": "Ubuntu", "release": "focal", "variant": "cloud"}
	Properties map[string]string `json:"properties" yaml:"properties"`

	// Whether the image is available to unauthenticated users
	// Example: false
	Public bool `json:"public" yaml:"public"`

	// When the image becomes obsolete
	// Example: 2025-03-23T20:00:00-04:00
	//
	// API extension: images_expiry
	ExpiresAt time.Time `json:"expires_at" yaml:"expires_at"`

	// List of profiles to use when creating from this image (if none provided by user)
	// Example: ["default"]
	//
	// API extension: image_profiles
	Profiles []string `json:"profiles" yaml:"profiles"`
}

// Image represents a LXD image
//
// swagger:model
type Image struct {
	ImagePut `yaml:",inline"`

	// List of aliases
	Aliases []ImageAlias `json:"aliases" yaml:"aliases"`

	// Architecture
	// Example: x86_64
	Architecture string `json:"architecture" yaml:"architecture"`

	// Whether the image is an automatically cached remote image
	// Example: true
	Cached bool `json:"cached" yaml:"cached"`

	// Original filename
	// Example: 06b86454720d36b20f94e31c6812e05ec51c1b568cf3a8abd273769d213394bb.rootfs
	Filename string `json:"filename" yaml:"filename"`

	// Full SHA-256 fingerprint
	// Example: 06b86454720d36b20f94e31c6812e05ec51c1b568cf3a8abd273769d213394bb
	Fingerprint string `json:"fingerprint" yaml:"fingerprint"`

	// Size of the image in bytes
	// Example: 272237676
	Size int64 `json:"size" yaml:"size"`

	// Where the image came from
	UpdateSource *ImageSource `json:"update_source,omitempty" yaml:"update_source,omitempty"`

	// Type of image (container or virtual-machine)
	// Example: container
	//
	// API extension: image_types
	Type string `json:"type" yaml:"type"`

	// When the image was originally created
	// Example: 2021-03-23T20:00:00-04:00
	CreatedAt time.Time `json:"created_at" yaml:"created_at"`

	// Last time the image was used
	// Example: 2021-03-22T20:39:00.575185384-04:00
	LastUsedAt time.Time `json:"last_used_at" yaml:"last_used_at"`

	// When the image was added to this LXD server
	// Example: 2021-03-24T14:18:15.115036787-04:00
	UploadedAt time.Time `json:"uploaded_at" yaml:"uploaded_at"`
}

// Writable converts a full Image struct into a ImagePut struct (filters read-only fields)
func (img *Image) Writable() ImagePut {
	return img.ImagePut
}

// ImageAlias represents an alias from the alias list of a LXD image
//
// swagger:model
type ImageAlias struct {
	// Name of the alias
	// Example: ubuntu-20.04
	Name string `json:"name" yaml:"name"`

	// Description of the alias
	// Example: Our preferred Ubuntu image
	Description string `json:"description" yaml:"description"`
}

// ImageSource represents the source of a LXD image
//
// swagger:model
type ImageSource struct {
	// Source alias to download from
	// Example: focal
	Alias string `json:"alias" yaml:"alias"`

	// Source server certificate (if not trusted by system CA)
	// Example: X509 PEM certificate
	Certificate string `json:"certificate" yaml:"certificate"`

	// Source server protocol
	// Example: simplestreams
	Protocol string `json:"protocol" yaml:"protocol"`

	// URL of the source server
	// Example: https://images.linuxcontainers.org
	Server string `json:"server" yaml:"server"`

	// Type of image (container or virtual-machine)
	// Example: container
	//
	// API extension: image_types
	ImageType string `json:"image_type" yaml:"image_type"`
}

// ImageAliasesPost represents a new LXD image alias
//
// swagger:model
type ImageAliasesPost struct {
	ImageAliasesEntry `yaml:",inline"`
}

// ImageAliasesEntryPost represents the required fields to rename a LXD image alias
//
// swagger:model
type ImageAliasesEntryPost struct {
	// Alias name
	// Example: ubuntu-20.04
	Name string `json:"name" yaml:"name"`
}

// ImageAliasesEntryPut represents the modifiable fields of a LXD image alias
//
// swagger:model
type ImageAliasesEntryPut struct {
	// Alias description
	// Example: Our preferred Ubuntu image
	Description string `json:"description" yaml:"description"`

	// Target fingerprint for the alias
	// Example: 06b86454720d36b20f94e31c6812e05ec51c1b568cf3a8abd273769d213394bb
	Target string `json:"target" yaml:"target"`
}

// ImageAliasesEntry represents a LXD image alias
//
// swagger:model
type ImageAliasesEntry struct {
	ImageAliasesEntryPut `yaml:",inline"`

	// Alias name
	// Example: ubuntu-20.04
	Name string `json:"name" yaml:"name"`

	// Alias type (container or virtual-machine)
	// Example: container
	//
	// API extension: image_types
	Type string `json:"type" yaml:"type"`
}

// ImageMetadata represents LXD image metadata (used in image tarball)
//
// swagger:model
type ImageMetadata struct {
	// Architecture name
	// Example: x86_64
	Architecture string `json:"architecture" yaml:"architecture"`

	// Image creation data (as UNIX epoch)
	// Example: 1620655439
	CreationDate int64 `json:"creation_date" yaml:"creation_date"`

	// Image expiry data (as UNIX epoch)
	// Example: 1620685757
	ExpiryDate int64 `json:"expiry_date" yaml:"expiry_date"`

	// Descriptive properties
	// Example: {"os": "Ubuntu", "release": "focal", "variant": "cloud"}
	Properties map[string]string `json:"properties" yaml:"properties"`

	// Template for files in the image
	Templates map[string]*ImageMetadataTemplate `json:"templates" yaml:"templates"`
}

// ImageMetadataTemplate represents a template entry in image metadata (used in image tarball)
//
// swagger:model
type ImageMetadataTemplate struct {
	// When to trigger the template (create, copy or start)
	// Example: create
	When []string `json:"when" yaml:"when"`

	// Whether to trigger only if the file is missing
	// Example: false
	CreateOnly bool `json:"create_only" yaml:"create_only"`

	// The template itself as a valid pongo2 template
	// Example: pongo2-template
	Template string `json:"template" yaml:"template"`

	// Key/value properties to pass to the template
	// Example: {"foo": "bar"}
	Properties map[string]string `json:"properties" yaml:"properties"`
}
<|endoftext|>


<commit_before>package httpd

import (
	"bufio"
	"fmt"
	"net"
	"net/http"

	"github.com/Symantec/Dominator/lib/format"
	"github.com/Symantec/Dominator/lib/url"
)

func (s state) listVMsHandler(w http.ResponseWriter, req *http.Request) {
	parsedQuery := url.ParseQuery(req.URL)
	writer := bufio.NewWriter(w)
	defer writer.Flush()
	ipAddrs := s.manager.ListVMs(true)
	matchState := parsedQuery.Table["state"]
	if parsedQuery.OutputType() == url.OutputTypeText && matchState == "" {
		for _, ipAddr := range ipAddrs {
			fmt.Fprintln(writer, ipAddr)
		}
		return
	}
	if parsedQuery.OutputType() == url.OutputTypeHtml {
		fmt.Fprintf(writer, "<title>List of VMs</title>\n")
		fmt.Fprintln(writer, ``)
		fmt.Fprintln(writer, "")
		fmt.Fprintln(writer, `<table border="1" style="width:100%">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    IP Addr	")
		fmt.Fprintln(writer, "    MAC Addr	")
		fmt.Fprintln(writer, "    Name(tag)	")
		fmt.Fprintln(writer, "    State	")
		fmt.Fprintln(writer, "    RAM	")
		fmt.Fprintln(writer, "    CPU	")
		fmt.Fprintln(writer, "    Num Volumes	")
		fmt.Fprintln(writer, "    Storage	")
		fmt.Fprintln(writer, "    Primary Owner	")
		fmt.Fprintln(writer, "
")
	}
	for _, ipAddr := range ipAddrs {
		vm, err := s.manager.GetVmInfo(net.ParseIP(ipAddr))
		if err != nil {
			continue
		}
		if matchState != "" && matchState != vm.State.String() {
			continue
		}
		switch parsedQuery.OutputType() {
		case url.OutputTypeText:
			fmt.Fprintln(writer, ipAddr)
		case url.OutputTypeHtml:
			if vm.Uncommitted {
				fmt.Fprintln(writer, "
")
			} else {
				fmt.Fprintln(writer, "
")
			}
			fmt.Fprintf(writer, "    %s	\n",
				ipAddr, ipAddr)
			fmt.Fprintf(writer, "    %s	\n", vm.Address.MacAddress)
			fmt.Fprintf(writer, "    %s	\n", vm.Tags["Name"])
			fmt.Fprintf(writer, "    %s	\n", vm.State)
			fmt.Fprintf(writer, "    %s	\n",
				format.FormatBytes(vm.MemoryInMiB<<20))
			fmt.Fprintf(writer, "    %g	\n",
				float64(vm.MilliCPUs)*1e-3)
			fmt.Fprintf(writer, "    %d	\n", len(vm.Volumes))
			var storage uint64
			for _, volume := range vm.Volumes {
				storage += volume.Size
			}
			fmt.Fprintf(writer, "    %s	\n",
				format.FormatBytes(storage))
			fmt.Fprintf(writer, "    %s	\n", vm.OwnerUsers[0])
			fmt.Fprintf(writer, "
\n")
		}
	}
	switch parsedQuery.OutputType() {
	case url.OutputTypeHtml:
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "")
	}
}
<commit_msg>Update hypervisor to report non-RAW volumes when listing VMs.<commit_after>package httpd

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"

	"github.com/Symantec/Dominator/lib/format"
	"github.com/Symantec/Dominator/lib/url"
	proto "github.com/Symantec/Dominator/proto/hypervisor"
)

func (s state) listVMsHandler(w http.ResponseWriter, req *http.Request) {
	parsedQuery := url.ParseQuery(req.URL)
	writer := bufio.NewWriter(w)
	defer writer.Flush()
	ipAddrs := s.manager.ListVMs(true)
	matchState := parsedQuery.Table["state"]
	if parsedQuery.OutputType() == url.OutputTypeText && matchState == "" {
		for _, ipAddr := range ipAddrs {
			fmt.Fprintln(writer, ipAddr)
		}
		return
	}
	if parsedQuery.OutputType() == url.OutputTypeHtml {
		fmt.Fprintf(writer, "<title>List of VMs</title>\n")
		fmt.Fprintln(writer, ``)
		fmt.Fprintln(writer, "")
		fmt.Fprintln(writer, `<table border="1" style="width:100%">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    IP Addr	")
		fmt.Fprintln(writer, "    MAC Addr	")
		fmt.Fprintln(writer, "    Name(tag)	")
		fmt.Fprintln(writer, "    State	")
		fmt.Fprintln(writer, "    RAM	")
		fmt.Fprintln(writer, "    CPU	")
		fmt.Fprintln(writer, "    Num Volumes	")
		fmt.Fprintln(writer, "    Storage	")
		fmt.Fprintln(writer, "    Primary Owner	")
		fmt.Fprintln(writer, "
")
	}
	for _, ipAddr := range ipAddrs {
		vm, err := s.manager.GetVmInfo(net.ParseIP(ipAddr))
		if err != nil {
			continue
		}
		if matchState != "" && matchState != vm.State.String() {
			continue
		}
		switch parsedQuery.OutputType() {
		case url.OutputTypeText:
			fmt.Fprintln(writer, ipAddr)
		case url.OutputTypeHtml:
			if vm.Uncommitted {
				fmt.Fprintln(writer, "
")
			} else {
				fmt.Fprintln(writer, "
")
			}
			fmt.Fprintf(writer, "    %s	\n",
				ipAddr, ipAddr)
			fmt.Fprintf(writer, "    %s	\n", vm.Address.MacAddress)
			fmt.Fprintf(writer, "    %s	\n", vm.Tags["Name"])
			fmt.Fprintf(writer, "    %s	\n", vm.State)
			fmt.Fprintf(writer, "    %s	\n",
				format.FormatBytes(vm.MemoryInMiB<<20))
			fmt.Fprintf(writer, "    %g	\n",
				float64(vm.MilliCPUs)*1e-3)
			writeNumVolumesTableEntry(writer, vm)
			writeStorageTotalTableEntry(writer, vm)
			fmt.Fprintf(writer, "    %s	\n", vm.OwnerUsers[0])
			fmt.Fprintf(writer, "
\n")
		}
	}
	switch parsedQuery.OutputType() {
	case url.OutputTypeHtml:
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "")
	}
}

func writeNumVolumesTableEntry(writer io.Writer, vm proto.VmInfo) {
	var comment string
	for _, volume := range vm.Volumes {
		if comment == "" && volume.Format != proto.VolumeFormatRaw {
			comment = `<font style="color:grey;font-size:12px"> (!RAW)</font>`
		}
	}
	fmt.Fprintf(writer, "    %d%s	\n", len(vm.Volumes), comment)
}

func writeStorageTotalTableEntry(writer io.Writer, vm proto.VmInfo) {
	var storage uint64
	for _, volume := range vm.Volumes {
		storage += volume.Size
	}
	fmt.Fprintf(writer, "    %s	\n", format.FormatBytes(storage))
}
<|endoftext|>


<commit_before>package control
<commit_msg>Add TestLoadPlugin test<commit_after>package control

import (
	"os"
	"path"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

var (
	PluginName = "pulse-collector-dummy"
	PulsePath  = os.Getenv("PULSE_PATH")
	PluginPath = path.Join(PulsePath, "plugin", "collector", PluginName)
)

// Uses the dummy collector plugin to simulate loading
func TestLoadPlugin(t *testing.T) {
	// These tests only work if PULSE_PATH is known
	// It is the responsibility of the testing framework to
	// build the plugins first into the build dir

	if PulsePath != "" {
		Convey("PluginManager.LoadPlugin", t, func() {

			Convey("loads plugin successfully", func() {
				p := PluginManager()
				p.Start()
				err := p.LoadPlugin(PluginPath)

				So(p.LoadedPlugins, ShouldNotBeEmpty)
				So(err, ShouldBeNil)
			})

			Convey("returns error if PluginManager is not started", func() {
				p := PluginManager()
				err := p.LoadPlugin(PluginPath)

				So(p.LoadedPlugins, ShouldBeEmpty)
				So(err, ShouldNotBeNil)
			})
		})

	}
}
<|endoftext|>


<commit_before>package allocrunner

import (
	"fmt"
	"time"

	multierror "github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
	clientconfig "github.com/hashicorp/nomad/client/config"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/client/taskenv"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
)

type hookResourceSetter interface {
	GetAllocHookResources() *cstructs.AllocHookResources
	SetAllocHookResources(*cstructs.AllocHookResources)
}

type allocHookResourceSetter struct {
	ar *allocRunner
}

func (a *allocHookResourceSetter) GetAllocHookResources() *cstructs.AllocHookResources {
	a.ar.hookStateMu.RLock()
	defer a.ar.hookStateMu.RUnlock()

	return a.ar.hookState
}

func (a *allocHookResourceSetter) SetAllocHookResources(res *cstructs.AllocHookResources) {
	a.ar.hookStateMu.Lock()
	defer a.ar.hookStateMu.Unlock()

	a.ar.hookState = res

	// Propagate to all of the TRs within the lock to ensure consistent state.
	// TODO: Refactor so TR's pull state from AR?
	for _, tr := range a.ar.tasks {
		tr.SetAllocHookResources(res)
	}
}

type networkIsolationSetter interface {
	SetNetworkIsolation(*drivers.NetworkIsolationSpec)
}

// allocNetworkIsolationSetter is a shim to allow the alloc network hook to
// set the alloc network isolation configuration without full access
// to the alloc runner
type allocNetworkIsolationSetter struct {
	ar *allocRunner
}

func (a *allocNetworkIsolationSetter) SetNetworkIsolation(n *drivers.NetworkIsolationSpec) {
	for _, tr := range a.ar.tasks {
		tr.SetNetworkIsolation(n)
	}
}

// allocHealthSetter is a shim to allow the alloc health watcher hook to set
// and clear the alloc health without full access to the alloc runner state
type allocHealthSetter struct {
	ar *allocRunner
}

// HasHealth returns true if a deployment status is already set.
func (a *allocHealthSetter) HasHealth() bool {
	a.ar.stateLock.Lock()
	defer a.ar.stateLock.Unlock()
	return a.ar.state.DeploymentStatus.HasHealth()
}

// ClearHealth allows the health watcher hook to clear the alloc's deployment
// health if the deployment id changes. It does not update the server as the
// status is only cleared when already receiving an update from the server.
//
// Only for use by health hook.
func (a *allocHealthSetter) ClearHealth() {
	a.ar.stateLock.Lock()
	a.ar.state.ClearDeploymentStatus()
	a.ar.persistDeploymentStatus(nil)
	a.ar.stateLock.Unlock()
}

// SetHealth allows the health watcher hook to set the alloc's
// deployment/migration health and emit task events.
//
// Only for use by health hook.
func (a *allocHealthSetter) SetHealth(healthy, isDeploy bool, trackerTaskEvents map[string]*structs.TaskEvent) {
	// Updating alloc deployment state is tricky because it may be nil, but
	// if it's not then we need to maintain the values of Canary and
	// ModifyIndex as they're only mutated by the server.
	a.ar.stateLock.Lock()
	a.ar.state.SetDeploymentStatus(time.Now(), healthy)
	a.ar.persistDeploymentStatus(a.ar.state.DeploymentStatus)
	terminalDesiredState := a.ar.Alloc().ServerTerminalStatus()
	a.ar.stateLock.Unlock()

	// If deployment is unhealthy emit task events explaining why
	if !healthy && isDeploy && !terminalDesiredState {
		for task, event := range trackerTaskEvents {
			if tr, ok := a.ar.tasks[task]; ok {
				// Append but don't emit event since the server
				// will be updated below
				tr.AppendEvent(event)
			}
		}
	}

	// Gather the state of the other tasks
	states := make(map[string]*structs.TaskState, len(a.ar.tasks))
	for name, tr := range a.ar.tasks {
		states[name] = tr.TaskState()
	}

	// Build the client allocation
	calloc := a.ar.clientAlloc(states)

	// Update the server
	a.ar.stateUpdater.AllocStateUpdated(calloc)

	// Broadcast client alloc to listeners
	a.ar.allocBroadcaster.Send(calloc)
}

// initRunnerHooks intializes the runners hooks.
func (ar *allocRunner) initRunnerHooks(config *clientconfig.Config) error {
	hookLogger := ar.logger.Named("runner_hook")

	// create health setting shim
	hs := &allocHealthSetter{ar}

	// create network isolation setting shim
	ns := &allocNetworkIsolationSetter{ar: ar}

	// create hook resource setting shim
	hrs := &allocHookResourceSetter{ar: ar}

	// build the network manager
	nm, err := newNetworkManager(ar.Alloc(), ar.driverManager)
	if err != nil {
		return fmt.Errorf("failed to configure network manager: %v", err)
	}

	// create network configurator
	nc, err := newNetworkConfigurator(hookLogger, ar.Alloc(), config)
	if err != nil {
		return fmt.Errorf("failed to initialize network configurator: %v", err)
	}

	// Create the alloc directory hook. This is run first to ensure the
	// directory path exists for other hooks.
	alloc := ar.Alloc()
	ar.runnerHooks = []interfaces.RunnerHook{
		newAllocDirHook(hookLogger, ar.allocDir),
		newUpstreamAllocsHook(hookLogger, ar.prevAllocWatcher),
		newDiskMigrationHook(hookLogger, ar.prevAllocMigrator, ar.allocDir),
		newAllocHealthWatcherHook(hookLogger, alloc, hs, ar.Listener(), ar.consulClient),
		newNetworkHook(hookLogger, ns, alloc, nm, nc),
		newGroupServiceHook(groupServiceHookConfig{
			alloc:          alloc,
			consul:         ar.consulClient,
			restarter:      ar,
			taskEnvBuilder: taskenv.NewBuilder(config.Node, ar.Alloc(), nil, config.Region).SetAllocDir(ar.allocDir.AllocDir),
			logger:         hookLogger,
		}),
		newConsulSockHook(hookLogger, alloc, ar.allocDir, config.ConsulConfig),
		newCSIHook(hookLogger, alloc, ar.rpcClient, ar.csiManager, hrs),
	}

	return nil
}

// prerun is used to run the runners prerun hooks.
func (ar *allocRunner) prerun() error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running pre-run hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished pre-run hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	for _, hook := range ar.runnerHooks {
		pre, ok := hook.(interfaces.RunnerPrerunHook)
		if !ok {
			continue
		}

		name := pre.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running pre-run hook", "name", name, "start", start)
		}

		if err := pre.Prerun(); err != nil {
			return fmt.Errorf("pre-run hook %q failed: %v", name, err)
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished pre-run hook", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return nil
}

// update runs the alloc runner update hooks. Update hooks are run
// asynchronously with all other alloc runner operations.
func (ar *allocRunner) update(update *structs.Allocation) error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running update hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished update hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	req := &interfaces.RunnerUpdateRequest{
		Alloc: update,
	}

	var merr multierror.Error
	for _, hook := range ar.runnerHooks {
		h, ok := hook.(interfaces.RunnerUpdateHook)
		if !ok {
			continue
		}

		name := h.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running update hook", "name", name, "start", start)
		}

		if err := h.Update(req); err != nil {
			merr.Errors = append(merr.Errors, fmt.Errorf("update hook %q failed: %v", name, err))
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished update hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return merr.ErrorOrNil()
}

// postrun is used to run the runners postrun hooks.
func (ar *allocRunner) postrun() error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running post-run hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished post-run hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	for _, hook := range ar.runnerHooks {
		post, ok := hook.(interfaces.RunnerPostrunHook)
		if !ok {
			continue
		}

		name := post.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running post-run hook", "name", name, "start", start)
		}

		if err := post.Postrun(); err != nil {
			return fmt.Errorf("hook %q failed: %v", name, err)
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished post-run hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return nil
}

// destroy is used to run the runners destroy hooks. All hooks are run and
// errors are returned as a multierror.
func (ar *allocRunner) destroy() error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running destroy hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished destroy hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	var merr multierror.Error
	for _, hook := range ar.runnerHooks {
		h, ok := hook.(interfaces.RunnerDestroyHook)
		if !ok {
			continue
		}

		name := h.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running destroy hook", "name", name, "start", start)
		}

		if err := h.Destroy(); err != nil {
			merr.Errors = append(merr.Errors, fmt.Errorf("destroy hook %q failed: %v", name, err))
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished destroy hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return merr.ErrorOrNil()
}

func (ar *allocRunner) preKillHooks() {
	for _, hook := range ar.runnerHooks {
		pre, ok := hook.(interfaces.RunnerPreKillHook)
		if !ok {
			continue
		}

		name := pre.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running alloc pre shutdown hook", "name", name, "start", start)
		}

		pre.PreKill()

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished alloc pre shutdown hook", "name", name, "end", end, "duration", end.Sub(start))
		}
	}
}

// shutdownHooks calls graceful shutdown hooks for when the agent is exiting.
func (ar *allocRunner) shutdownHooks() {
	for _, hook := range ar.runnerHooks {
		sh, ok := hook.(interfaces.ShutdownHook)
		if !ok {
			continue
		}

		name := sh.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running shutdown hook", "name", name, "start", start)
		}

		sh.Shutdown()

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished shutdown hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}
}
<commit_msg>hook resources: Init with empty resources during setup<commit_after>package allocrunner

import (
	"fmt"
	"time"

	multierror "github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
	clientconfig "github.com/hashicorp/nomad/client/config"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/client/taskenv"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
)

type hookResourceSetter interface {
	GetAllocHookResources() *cstructs.AllocHookResources
	SetAllocHookResources(*cstructs.AllocHookResources)
}

type allocHookResourceSetter struct {
	ar *allocRunner
}

func (a *allocHookResourceSetter) GetAllocHookResources() *cstructs.AllocHookResources {
	a.ar.hookStateMu.RLock()
	defer a.ar.hookStateMu.RUnlock()

	return a.ar.hookState
}

func (a *allocHookResourceSetter) SetAllocHookResources(res *cstructs.AllocHookResources) {
	a.ar.hookStateMu.Lock()
	defer a.ar.hookStateMu.Unlock()

	a.ar.hookState = res

	// Propagate to all of the TRs within the lock to ensure consistent state.
	// TODO: Refactor so TR's pull state from AR?
	for _, tr := range a.ar.tasks {
		tr.SetAllocHookResources(res)
	}
}

type networkIsolationSetter interface {
	SetNetworkIsolation(*drivers.NetworkIsolationSpec)
}

// allocNetworkIsolationSetter is a shim to allow the alloc network hook to
// set the alloc network isolation configuration without full access
// to the alloc runner
type allocNetworkIsolationSetter struct {
	ar *allocRunner
}

func (a *allocNetworkIsolationSetter) SetNetworkIsolation(n *drivers.NetworkIsolationSpec) {
	for _, tr := range a.ar.tasks {
		tr.SetNetworkIsolation(n)
	}
}

// allocHealthSetter is a shim to allow the alloc health watcher hook to set
// and clear the alloc health without full access to the alloc runner state
type allocHealthSetter struct {
	ar *allocRunner
}

// HasHealth returns true if a deployment status is already set.
func (a *allocHealthSetter) HasHealth() bool {
	a.ar.stateLock.Lock()
	defer a.ar.stateLock.Unlock()
	return a.ar.state.DeploymentStatus.HasHealth()
}

// ClearHealth allows the health watcher hook to clear the alloc's deployment
// health if the deployment id changes. It does not update the server as the
// status is only cleared when already receiving an update from the server.
//
// Only for use by health hook.
func (a *allocHealthSetter) ClearHealth() {
	a.ar.stateLock.Lock()
	a.ar.state.ClearDeploymentStatus()
	a.ar.persistDeploymentStatus(nil)
	a.ar.stateLock.Unlock()
}

// SetHealth allows the health watcher hook to set the alloc's
// deployment/migration health and emit task events.
//
// Only for use by health hook.
func (a *allocHealthSetter) SetHealth(healthy, isDeploy bool, trackerTaskEvents map[string]*structs.TaskEvent) {
	// Updating alloc deployment state is tricky because it may be nil, but
	// if it's not then we need to maintain the values of Canary and
	// ModifyIndex as they're only mutated by the server.
	a.ar.stateLock.Lock()
	a.ar.state.SetDeploymentStatus(time.Now(), healthy)
	a.ar.persistDeploymentStatus(a.ar.state.DeploymentStatus)
	terminalDesiredState := a.ar.Alloc().ServerTerminalStatus()
	a.ar.stateLock.Unlock()

	// If deployment is unhealthy emit task events explaining why
	if !healthy && isDeploy && !terminalDesiredState {
		for task, event := range trackerTaskEvents {
			if tr, ok := a.ar.tasks[task]; ok {
				// Append but don't emit event since the server
				// will be updated below
				tr.AppendEvent(event)
			}
		}
	}

	// Gather the state of the other tasks
	states := make(map[string]*structs.TaskState, len(a.ar.tasks))
	for name, tr := range a.ar.tasks {
		states[name] = tr.TaskState()
	}

	// Build the client allocation
	calloc := a.ar.clientAlloc(states)

	// Update the server
	a.ar.stateUpdater.AllocStateUpdated(calloc)

	// Broadcast client alloc to listeners
	a.ar.allocBroadcaster.Send(calloc)
}

// initRunnerHooks intializes the runners hooks.
func (ar *allocRunner) initRunnerHooks(config *clientconfig.Config) error {
	hookLogger := ar.logger.Named("runner_hook")

	// create health setting shim
	hs := &allocHealthSetter{ar}

	// create network isolation setting shim
	ns := &allocNetworkIsolationSetter{ar: ar}

	// create hook resource setting shim
	hrs := &allocHookResourceSetter{ar: ar}
	hrs.SetAllocHookResources(&cstructs.AllocHookResources{})

	// build the network manager
	nm, err := newNetworkManager(ar.Alloc(), ar.driverManager)
	if err != nil {
		return fmt.Errorf("failed to configure network manager: %v", err)
	}

	// create network configurator
	nc, err := newNetworkConfigurator(hookLogger, ar.Alloc(), config)
	if err != nil {
		return fmt.Errorf("failed to initialize network configurator: %v", err)
	}

	// Create the alloc directory hook. This is run first to ensure the
	// directory path exists for other hooks.
	alloc := ar.Alloc()
	ar.runnerHooks = []interfaces.RunnerHook{
		newAllocDirHook(hookLogger, ar.allocDir),
		newUpstreamAllocsHook(hookLogger, ar.prevAllocWatcher),
		newDiskMigrationHook(hookLogger, ar.prevAllocMigrator, ar.allocDir),
		newAllocHealthWatcherHook(hookLogger, alloc, hs, ar.Listener(), ar.consulClient),
		newNetworkHook(hookLogger, ns, alloc, nm, nc),
		newGroupServiceHook(groupServiceHookConfig{
			alloc:          alloc,
			consul:         ar.consulClient,
			restarter:      ar,
			taskEnvBuilder: taskenv.NewBuilder(config.Node, ar.Alloc(), nil, config.Region).SetAllocDir(ar.allocDir.AllocDir),
			logger:         hookLogger,
		}),
		newConsulSockHook(hookLogger, alloc, ar.allocDir, config.ConsulConfig),
		newCSIHook(hookLogger, alloc, ar.rpcClient, ar.csiManager, hrs),
	}

	return nil
}

// prerun is used to run the runners prerun hooks.
func (ar *allocRunner) prerun() error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running pre-run hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished pre-run hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	for _, hook := range ar.runnerHooks {
		pre, ok := hook.(interfaces.RunnerPrerunHook)
		if !ok {
			continue
		}

		name := pre.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running pre-run hook", "name", name, "start", start)
		}

		if err := pre.Prerun(); err != nil {
			return fmt.Errorf("pre-run hook %q failed: %v", name, err)
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished pre-run hook", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return nil
}

// update runs the alloc runner update hooks. Update hooks are run
// asynchronously with all other alloc runner operations.
func (ar *allocRunner) update(update *structs.Allocation) error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running update hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished update hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	req := &interfaces.RunnerUpdateRequest{
		Alloc: update,
	}

	var merr multierror.Error
	for _, hook := range ar.runnerHooks {
		h, ok := hook.(interfaces.RunnerUpdateHook)
		if !ok {
			continue
		}

		name := h.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running update hook", "name", name, "start", start)
		}

		if err := h.Update(req); err != nil {
			merr.Errors = append(merr.Errors, fmt.Errorf("update hook %q failed: %v", name, err))
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished update hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return merr.ErrorOrNil()
}

// postrun is used to run the runners postrun hooks.
func (ar *allocRunner) postrun() error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running post-run hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished post-run hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	for _, hook := range ar.runnerHooks {
		post, ok := hook.(interfaces.RunnerPostrunHook)
		if !ok {
			continue
		}

		name := post.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running post-run hook", "name", name, "start", start)
		}

		if err := post.Postrun(); err != nil {
			return fmt.Errorf("hook %q failed: %v", name, err)
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished post-run hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return nil
}

// destroy is used to run the runners destroy hooks. All hooks are run and
// errors are returned as a multierror.
func (ar *allocRunner) destroy() error {
	if ar.logger.IsTrace() {
		start := time.Now()
		ar.logger.Trace("running destroy hooks", "start", start)
		defer func() {
			end := time.Now()
			ar.logger.Trace("finished destroy hooks", "end", end, "duration", end.Sub(start))
		}()
	}

	var merr multierror.Error
	for _, hook := range ar.runnerHooks {
		h, ok := hook.(interfaces.RunnerDestroyHook)
		if !ok {
			continue
		}

		name := h.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running destroy hook", "name", name, "start", start)
		}

		if err := h.Destroy(); err != nil {
			merr.Errors = append(merr.Errors, fmt.Errorf("destroy hook %q failed: %v", name, err))
		}

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished destroy hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}

	return merr.ErrorOrNil()
}

func (ar *allocRunner) preKillHooks() {
	for _, hook := range ar.runnerHooks {
		pre, ok := hook.(interfaces.RunnerPreKillHook)
		if !ok {
			continue
		}

		name := pre.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running alloc pre shutdown hook", "name", name, "start", start)
		}

		pre.PreKill()

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished alloc pre shutdown hook", "name", name, "end", end, "duration", end.Sub(start))
		}
	}
}

// shutdownHooks calls graceful shutdown hooks for when the agent is exiting.
func (ar *allocRunner) shutdownHooks() {
	for _, hook := range ar.runnerHooks {
		sh, ok := hook.(interfaces.ShutdownHook)
		if !ok {
			continue
		}

		name := sh.Name()
		var start time.Time
		if ar.logger.IsTrace() {
			start = time.Now()
			ar.logger.Trace("running shutdown hook", "name", name, "start", start)
		}

		sh.Shutdown()

		if ar.logger.IsTrace() {
			end := time.Now()
			ar.logger.Trace("finished shutdown hooks", "name", name, "end", end, "duration", end.Sub(start))
		}
	}
}
<|endoftext|>


<commit_before>package honeybadger

import "fmt"

func newBufferedWorker() Worker {
	worker := make(BufferedWorker, 100)
	go func() {
		for w := range worker {
			work := func() error {
				defer func() {
					if err := recover(); err != nil {
						fmt.Printf("worker recovered from panic: %v\n", err)
					}
				}()
				return w()
			}
			if err := work(); err != nil {
				fmt.Printf("worker processing error: %v\n", err)
			}
		}
	}()
	return worker
}

type BufferedWorker chan Envelope

func (worker BufferedWorker) Push(work Envelope) error {
	select {
	case worker <- work:
		return nil
	default:
		return fmt.Errorf("the channel is full.")
	}
}

func (worker BufferedWorker) Flush() error {
	ch := make(chan bool)
	worker <- func() error {
		ch <- true
		return nil
	}
	<-ch
	return nil
}
<commit_msg>Refactor error.<commit_after>package honeybadger

import "fmt"

var (
	WorkerOverflowError = fmt.Errorf("The worker is full; this envelope will be dropped.")
)

func newBufferedWorker() Worker {
	worker := make(BufferedWorker, 100)
	go func() {
		for w := range worker {
			work := func() error {
				defer func() {
					if err := recover(); err != nil {
						fmt.Printf("worker recovered from panic: %v\n", err)
					}
				}()
				return w()
			}
			if err := work(); err != nil {
				fmt.Printf("worker processing error: %v\n", err)
			}
		}
	}()
	return worker
}

type BufferedWorker chan Envelope

func (worker BufferedWorker) Push(work Envelope) error {
	select {
	case worker <- work:
		return nil
	default:
		return WorkerOverflowError
	}
}

func (worker BufferedWorker) Flush() error {
	ch := make(chan bool)
	worker <- func() error {
		ch <- true
		return nil
	}
	<-ch
	return nil
}
<|endoftext|>


<commit_before>package bufpool

import (
	"bytes"
	"sync"
)

// AverageBufferSize should be adjusted to the average size of a bytes.buffer
// in your application.
var AverageBufferSize int = 4096

var bufferPool = &sync.Pool{
	New: func() interface{} {
		b := bytes.NewBuffer(make([]byte, AverageBufferSize))
		b.Reset()
		return b
	},
}

// Get returns a buffer from the pool.
func Get() *bytes.Buffer {
	return bufferPool.Get().(*bytes.Buffer)
}

// Put returns a buffer to the pool.
// The buffer is reset before it is put back into circulation.
func Put(buf *bytes.Buffer) {
	buf.Reset()
	bufferPool.Put(buf)
}
<commit_msg>fix golint waring<commit_after>package bufpool

import (
	"bytes"
	"sync"
)

// AverageBufferSize should be adjusted to the average size of a bytes.buffer
// in your application.
const AverageBufferSize = 4096

var bufferPool = &sync.Pool{
	New: func() interface{} {
		b := bytes.NewBuffer(make([]byte, AverageBufferSize))
		b.Reset()
		return b
	},
}

// Get returns a buffer from the pool.
func Get() *bytes.Buffer {
	return bufferPool.Get().(*bytes.Buffer)
}

// Put returns a buffer to the pool.
// The buffer is reset before it is put back into circulation.
func Put(buf *bytes.Buffer) {
	buf.Reset()
	bufferPool.Put(buf)
}
<|endoftext|>