<commit_before>package main

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"path"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/docker/swarm/api"
	"github.com/docker/swarm/cluster"
	"github.com/docker/swarm/discovery"
	"github.com/docker/swarm/scheduler"
	"github.com/docker/swarm/scheduler/filter"
	"github.com/docker/swarm/scheduler/strategy"
	"github.com/docker/swarm/state"
)

type logHandler struct {
}

func (h *logHandler) Handle(e *cluster.Event) error {
	log.Printf("event -> status: %q from: %q id: %q node: %q", e.Status, e.From, e.Id, e.Node.Name)
	return nil
}

// Load the TLS certificates/keys and, if verify is true, the CA.
func loadTlsConfig(ca, cert, key string, verify bool) (*tls.Config, error) {
	c, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, fmt.Errorf("Couldn't load X509 key pair (%s, %s): %s. Key encrypted?",
			cert, key, err)
	}

	config := &tls.Config{
		Certificates: []tls.Certificate{c},
		MinVersion:   tls.VersionTLS10,
	}

	if verify {
		certPool := x509.NewCertPool()
		file, err := ioutil.ReadFile(ca)
		if err != nil {
			return nil, fmt.Errorf("Couldn't read CA certificate: %s", err)
		}
		certPool.AppendCertsFromPEM(file)
		config.RootCAs = certPool
	} else {
		// If --tlsverify is not supplied, disable CA validation.
		config.InsecureSkipVerify = true
	}

	return config, nil
}

func manage(c *cli.Context) {
	var (
		tlsConfig *tls.Config = nil
		err       error
	)

	// If either --tls or --tlsverify are specified, load the certificates.
	if c.Bool("tls") || c.Bool("tlsverify") {
		tlsConfig, err = loadTlsConfig(
			c.String("tlscacert"),
			c.String("tlscert"),
			c.String("tlskey"),
			c.Bool("tlsverify"))
		if err != nil {
			log.Fatal(err)
		}
	}

	store := state.NewStore(path.Join(c.String("rootdir"), "state"))
	if err := store.Initialize(); err != nil {
		log.Fatal(err)
	}

	cluster := cluster.NewCluster(store, tlsConfig, c.Float64("overcommit"))
	cluster.Events(&logHandler{})

	if !c.IsSet("discovery") {
		log.Fatal("--discovery required to manage a cluster")
	}

	s, err := strategy.New(c.String("strategy"))
	if err != nil {
		log.Fatal(err)
	}

	// see https://github.com/codegangsta/cli/issues/160
	names := c.StringSlice("filter")
	if c.IsSet("filter") || c.IsSet("f") {
		names = names[3:]
	}
	fs, err := filter.New(names)
	if err != nil {
		log.Fatal(err)
	}

	// get the list of nodes from the discovery service
	go func() {
		d, err := discovery.New(c.String("discovery"), c.Int("heartbeat"))
		if err != nil {
			log.Fatal(err)
		}

		nodes, err := d.Fetch()
		if err != nil {
			log.Fatal(err)

		}
		cluster.UpdateNodes(nodes)

		go d.Watch(cluster.UpdateNodes)
	}()

	sched := scheduler.NewScheduler(
		cluster,
		s,
		fs,
	)

	// see https://github.com/codegangsta/cli/issues/160
	hosts := c.StringSlice("host")
	if c.IsSet("host") || c.IsSet("H") {
		hosts = hosts[1:]
	}
	log.Fatal(api.ListenAndServe(cluster, sched, hosts, c.App.Version, c.Bool("cors"), tlsConfig))
}
<commit_msg>fix SWARM_DISCOVERY env var<commit_after>package main

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"path"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/docker/swarm/api"
	"github.com/docker/swarm/cluster"
	"github.com/docker/swarm/discovery"
	"github.com/docker/swarm/scheduler"
	"github.com/docker/swarm/scheduler/filter"
	"github.com/docker/swarm/scheduler/strategy"
	"github.com/docker/swarm/state"
)

type logHandler struct {
}

func (h *logHandler) Handle(e *cluster.Event) error {
	log.Printf("event -> status: %q from: %q id: %q node: %q", e.Status, e.From, e.Id, e.Node.Name)
	return nil
}

// Load the TLS certificates/keys and, if verify is true, the CA.
func loadTlsConfig(ca, cert, key string, verify bool) (*tls.Config, error) {
	c, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, fmt.Errorf("Couldn't load X509 key pair (%s, %s): %s. Key encrypted?",
			cert, key, err)
	}

	config := &tls.Config{
		Certificates: []tls.Certificate{c},
		MinVersion:   tls.VersionTLS10,
	}

	if verify {
		certPool := x509.NewCertPool()
		file, err := ioutil.ReadFile(ca)
		if err != nil {
			return nil, fmt.Errorf("Couldn't read CA certificate: %s", err)
		}
		certPool.AppendCertsFromPEM(file)
		config.RootCAs = certPool
	} else {
		// If --tlsverify is not supplied, disable CA validation.
		config.InsecureSkipVerify = true
	}

	return config, nil
}

func manage(c *cli.Context) {
	var (
		tlsConfig *tls.Config = nil
		err       error
	)

	// If either --tls or --tlsverify are specified, load the certificates.
	if c.Bool("tls") || c.Bool("tlsverify") {
		tlsConfig, err = loadTlsConfig(
			c.String("tlscacert"),
			c.String("tlscert"),
			c.String("tlskey"),
			c.Bool("tlsverify"))
		if err != nil {
			log.Fatal(err)
		}
	}

	store := state.NewStore(path.Join(c.String("rootdir"), "state"))
	if err := store.Initialize(); err != nil {
		log.Fatal(err)
	}

	cluster := cluster.NewCluster(store, tlsConfig, c.Float64("overcommit"))
	cluster.Events(&logHandler{})

	if c.String("discovery") == "" {
		log.Fatal("--discovery required to manage a cluster")
	}

	s, err := strategy.New(c.String("strategy"))
	if err != nil {
		log.Fatal(err)
	}

	// see https://github.com/codegangsta/cli/issues/160
	names := c.StringSlice("filter")
	if c.IsSet("filter") || c.IsSet("f") {
		names = names[3:]
	}
	fs, err := filter.New(names)
	if err != nil {
		log.Fatal(err)
	}

	// get the list of nodes from the discovery service
	go func() {
		d, err := discovery.New(c.String("discovery"), c.Int("heartbeat"))
		if err != nil {
			log.Fatal(err)
		}

		nodes, err := d.Fetch()
		if err != nil {
			log.Fatal(err)

		}
		cluster.UpdateNodes(nodes)

		go d.Watch(cluster.UpdateNodes)
	}()

	sched := scheduler.NewScheduler(
		cluster,
		s,
		fs,
	)

	// see https://github.com/codegangsta/cli/issues/160
	hosts := c.StringSlice("host")
	if c.IsSet("host") || c.IsSet("H") {
		hosts = hosts[1:]
	}
	log.Fatal(api.ListenAndServe(cluster, sched, hosts, c.App.Version, c.Bool("cors"), tlsConfig))
}
<|endoftext|>


<commit_before>// Package delta provides a log handler which times the delta
// between each log call, useful for debug output for command-line
// programs.
package delta

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/apex/log"
	"github.com/aybabtme/rgbterm"
	"github.com/tj/go-spin"
)

// TODO: move colors and share in text handler etc

// color function.
type colorFunc func(string) string

// gray string.
func gray(s string) string {
	return rgbterm.FgString(s, 150, 150, 150)
}

// blue string.
func blue(s string) string {
	return rgbterm.FgString(s, 77, 173, 247)
}

// cyan string.
func cyan(s string) string {
	return rgbterm.FgString(s, 34, 184, 207)
}

// green string.
func green(s string) string {
	return rgbterm.FgString(s, 0, 200, 255)
}

// red string.
func red(s string) string {
	return rgbterm.FgString(s, 194, 37, 92)
}

// yellow string.
func yellow(s string) string {
	return rgbterm.FgString(s, 252, 196, 25)
}

// Colors mapping.
var Colors = [...]colorFunc{
	log.DebugLevel: gray,
	log.InfoLevel:  blue,
	log.WarnLevel:  yellow,
	log.ErrorLevel: red,
	log.FatalLevel: red,
}

// Strings mapping.
var Strings = [...]string{
	log.DebugLevel: "DEBU",
	log.InfoLevel:  "INFO",
	log.WarnLevel:  "WARN",
	log.ErrorLevel: "ERRO",
	log.FatalLevel: "FATA",
}

// Default handler.
var Default = New(os.Stderr)

// Handler implementation.
type Handler struct {
	entries chan *log.Entry
	start   time.Time
	spin    *spin.Spinner
	prev    *log.Entry
	done    chan struct{}
	w       io.Writer
}

// New handler.
func New(w io.Writer) *Handler {
	h := &Handler{
		entries: make(chan *log.Entry),
		done:    make(chan struct{}),
		start:   time.Now(),
		spin:    spin.New(),
		w:       w,
	}

	go h.loop()

	return h
}

// Close the handler.
func (h *Handler) Close() error {
	h.done <- struct{}{}
	close(h.done)
	close(h.entries)
	return nil
}

// loop for rendering.
func (h *Handler) loop() {
	ticker := time.NewTicker(100 * time.Millisecond)

	for {
		select {
		case e := <-h.entries:
			if h.prev != nil {
				h.render(h.prev, true)
			}
			h.render(e, false)
			h.prev = e
		case <-ticker.C:
			if h.prev != nil {
				h.render(h.prev, false)
			}
			h.spin.Next()
		case <-h.done:
			ticker.Stop()
			if h.prev != nil {
				h.render(h.prev, true)
			}
			return
		}
	}
}

func (h *Handler) render(e *log.Entry, done bool) {
	color := Colors[e.Level]
	level := Strings[e.Level]
	names := e.Fields.Names()

	// delta and spinner
	if done {
		fmt.Fprintf(h.w, "\r     %7s", time.Since(h.start).Round(time.Millisecond))
	} else {
		fmt.Fprintf(h.w, "\r   %s %7s", h.spin.Current(), time.Since(h.start).Round(time.Millisecond))
	}

	// message
	fmt.Fprintf(h.w, " %s %s", color(level), color(e.Message))

	// fields
	for _, name := range names {
		v := e.Fields.Get(name)

		if v == "" {
			continue
		}

		fmt.Fprintf(h.w, " %s%s%v", color(name), gray("="), v)
	}

	// newline
	if done {
		fmt.Fprintf(h.w, "\n")
		h.start = time.Now()
	}
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	h.entries <- e
	return nil
}
<commit_msg>change delta to left-align<commit_after>// Package delta provides a log handler which times the delta
// between each log call, useful for debug output for command-line
// programs.
package delta

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/apex/log"
	"github.com/aybabtme/rgbterm"
	"github.com/tj/go-spin"
)

// TODO: move colors and share in text handler etc

// color function.
type colorFunc func(string) string

// gray string.
func gray(s string) string {
	return rgbterm.FgString(s, 150, 150, 150)
}

// blue string.
func blue(s string) string {
	return rgbterm.FgString(s, 77, 173, 247)
}

// cyan string.
func cyan(s string) string {
	return rgbterm.FgString(s, 34, 184, 207)
}

// green string.
func green(s string) string {
	return rgbterm.FgString(s, 0, 200, 255)
}

// red string.
func red(s string) string {
	return rgbterm.FgString(s, 194, 37, 92)
}

// yellow string.
func yellow(s string) string {
	return rgbterm.FgString(s, 252, 196, 25)
}

// Colors mapping.
var Colors = [...]colorFunc{
	log.DebugLevel: gray,
	log.InfoLevel:  blue,
	log.WarnLevel:  yellow,
	log.ErrorLevel: red,
	log.FatalLevel: red,
}

// Strings mapping.
var Strings = [...]string{
	log.DebugLevel: "DEBU",
	log.InfoLevel:  "INFO",
	log.WarnLevel:  "WARN",
	log.ErrorLevel: "ERRO",
	log.FatalLevel: "FATA",
}

// Default handler.
var Default = New(os.Stderr)

// Handler implementation.
type Handler struct {
	entries chan *log.Entry
	start   time.Time
	spin    *spin.Spinner
	prev    *log.Entry
	done    chan struct{}
	w       io.Writer
}

// New handler.
func New(w io.Writer) *Handler {
	h := &Handler{
		entries: make(chan *log.Entry),
		done:    make(chan struct{}),
		start:   time.Now(),
		spin:    spin.New(),
		w:       w,
	}

	go h.loop()

	return h
}

// Close the handler.
func (h *Handler) Close() error {
	h.done <- struct{}{}
	close(h.done)
	close(h.entries)
	return nil
}

// loop for rendering.
func (h *Handler) loop() {
	ticker := time.NewTicker(100 * time.Millisecond)

	for {
		select {
		case e := <-h.entries:
			if h.prev != nil {
				h.render(h.prev, true)
			}
			h.render(e, false)
			h.prev = e
		case <-ticker.C:
			if h.prev != nil {
				h.render(h.prev, false)
			}
			h.spin.Next()
		case <-h.done:
			ticker.Stop()
			if h.prev != nil {
				h.render(h.prev, true)
			}
			return
		}
	}
}

func (h *Handler) render(e *log.Entry, done bool) {
	color := Colors[e.Level]
	level := Strings[e.Level]
	names := e.Fields.Names()

	// delta and spinner
	if done {
		fmt.Fprintf(h.w, "\r     %-7s", time.Since(h.start).Round(time.Millisecond))
	} else {
		fmt.Fprintf(h.w, "\r   %s %-7s", h.spin.Current(), time.Since(h.start).Round(time.Millisecond))
	}

	// message
	fmt.Fprintf(h.w, " %s %s", color(level), color(e.Message))

	// fields
	for _, name := range names {
		v := e.Fields.Get(name)

		if v == "" {
			continue
		}

		fmt.Fprintf(h.w, " %s%s%v", color(name), gray("="), v)
	}

	// newline
	if done {
		fmt.Fprintf(h.w, "\n")
		h.start = time.Now()
	}
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	h.entries <- e
	return nil
}
<|endoftext|>


<commit_before>package shadow

import (
	"fmt"
	"io"
	"reflect"

	"github.com/hashicorp/go-multierror"
	"github.com/mitchellh/reflectwalk"
)

// Close will close all shadow values within the given structure.
//
// This uses reflection to walk the structure, find all shadow elements,
// and close them. Currently this will only find struct fields that are
// shadow values, and not slice elements, etc.
func Close(v interface{}) error {
	// We require a pointer so we can address the internal fields
	val := reflect.ValueOf(v)
	if val.Kind() != reflect.Ptr {
		return fmt.Errorf("value must be a pointer")
	}

	// Walk and close
	var w closeWalker
	if err := reflectwalk.Walk(v, &w); err != nil {
		return err
	}

	return w.Err
}

type closeWalker struct {
	Err error
}

func (w *closeWalker) Struct(reflect.Value) error {
	// Do nothing. We implement this for reflectwalk.StructWalker
	return nil
}

func (w *closeWalker) StructField(f reflect.StructField, v reflect.Value) error {
	// Not sure why this would be but lets avoid some panics
	if !v.IsValid() {
		return nil
	}

	// Empty for exported, so don't check unexported fields
	if f.PkgPath != "" {
		return nil
	}

	// Verify the io.Closer is in this package
	typ := v.Type()
	if typ.PkgPath() != "github.com/hashicorp/terraform/helper/shadow" {
		return nil
	}

	// We're looking for an io.Closer
	raw := v.Interface()
	if raw == nil {
		return nil
	}

	closer, ok := raw.(io.Closer)
	if !ok && v.CanAddr() {
		closer, ok = v.Addr().Interface().(io.Closer)
	}
	if !ok {
		return reflectwalk.SkipEntry
	}

	// Close it
	if err := closer.Close(); err != nil {
		w.Err = multierror.Append(w.Err, err)
	}

	// Don't go into the struct field
	return reflectwalk.SkipEntry
}
<commit_msg>Make sure shadow.closeWalker doesn't copy Mutexes<commit_after>package shadow

import (
	"fmt"
	"io"
	"reflect"

	"github.com/hashicorp/go-multierror"
	"github.com/mitchellh/reflectwalk"
)

// Close will close all shadow values within the given structure.
//
// This uses reflection to walk the structure, find all shadow elements,
// and close them. Currently this will only find struct fields that are
// shadow values, and not slice elements, etc.
func Close(v interface{}) error {
	// We require a pointer so we can address the internal fields
	val := reflect.ValueOf(v)
	if val.Kind() != reflect.Ptr {
		return fmt.Errorf("value must be a pointer")
	}

	// Walk and close
	var w closeWalker
	if err := reflectwalk.Walk(v, &w); err != nil {
		return err
	}

	return w.Err
}

type closeWalker struct {
	Err error
}

func (w *closeWalker) Struct(reflect.Value) error {
	// Do nothing. We implement this for reflectwalk.StructWalker
	return nil
}

var closerType = reflect.TypeOf((*io.Closer)(nil)).Elem()

func (w *closeWalker) StructField(f reflect.StructField, v reflect.Value) error {
	// Not sure why this would be but lets avoid some panics
	if !v.IsValid() {
		return nil
	}

	// Empty for exported, so don't check unexported fields
	if f.PkgPath != "" {
		return nil
	}

	// Verify the io.Closer is in this package
	typ := v.Type()
	if typ.PkgPath() != "github.com/hashicorp/terraform/helper/shadow" {
		return nil
	}

	var closer io.Closer
	if v.Type().Implements(closerType) {
		closer = v.Interface().(io.Closer)
	} else if v.CanAddr() {
		// The Close method may require a pointer receiver, but we only have a value.
		v := v.Addr()
		if v.Type().Implements(closerType) {
			closer = v.Interface().(io.Closer)
		}
	}

	if closer == nil {
		return reflectwalk.SkipEntry
	}

	// Close it
	if err := closer.Close(); err != nil {
		w.Err = multierror.Append(w.Err, err)
	}

	// Don't go into the struct field
	return reflectwalk.SkipEntry
}
<|endoftext|>


<commit_before>package main

import (
    "encoding/json"
    "fmt"
    "io/ioutil"
    "regexp"
    "strconv"
    "os"
    "jfb/svgxml"
)

type Config struct {
    // member  type        json element
    Colours map[string]string `json:"colours"`
}

func main() {

    var config Config
    var colors map[int]string

    if len(os.Args) < 2 {
        fmt.Fprintf(os.Stderr, "Please specify at least one filename on the command line.\n")
        os.Exit(1)
    }

    jsoncfg, err := ioutil.ReadFile("config.json")
    if err != nil {
        panic(err)
    }

    mapsvg, err := ioutil.ReadFile(os.Args[1])
    if err != nil {
        panic(err)
    }

    err = json.Unmarshal(jsoncfg, &config)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Unmarshal: ")
        panic(err)
    }

    // convert string-based JSON to all-integer map
    colors = make(map[int]string)
    for k, v := range config.Colours {
        k_i, _ := strconv.ParseInt(k, 0, 64)
        colors[int(k_i)] = v
    }

    mapsvg_obj := svgxml.XML2SVG(mapsvg)
    if mapsvg_obj == nil {
        os.Exit(1)
    }

    re_fill, err := regexp.Compile(`(fill:#)......`)
    if err != nil {
        panic(err)
    }

    element := svgxml.FindPathById(mapsvg_obj, "HI")

    if element != nil {
        element.Style = string(re_fill.ReplaceAll([]byte(element.Style), []byte("${1}ffaaaa")))
    } else {
        fmt.Println("'HI' not found")
    }

    // fmt.Println(string(svgxml.SVG2XML(mapsvg_obj, true)))

}

// vim:ai:sw=4:ts=8:
<commit_msg>colouring of (SVG) maps based on (go) map data<commit_after>package main

import (
    "encoding/json"
    "fmt"
    "io/ioutil"
    "regexp"
    "strconv"
    "os"
    "sort"
    "jfb/svgxml"
)

type Config struct {
    Colours map[string]string `json:"colours"`
}

// set up integer array sorting

type IntArray []int

func (list IntArray) Len() int {
    return len(list)
}

func (list IntArray) Swap(a, b int) {
    list[a], list[b] = list[b], list[a]
}

func (list IntArray) Less(a, b int) bool {
    return list[a] < list[b]
}

func main() {

    var config Config

    if len(os.Args) < 2 {
        fmt.Fprintf(os.Stderr, "Please specify at least one filename on the command line.\n")
        os.Exit(1)
    }

    jsoncfg, err := ioutil.ReadFile("config.json")
    if err != nil {
        panic(err)
    }

    mapsvg, err := ioutil.ReadFile(os.Args[1])
    if err != nil {
        panic(err)
    }

    err = json.Unmarshal(jsoncfg, &config)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Unmarshal: ")
        panic(err)
    }

    // make sorted list of keys (minimum counts) for later comparisons
    mincount := make([]int, len(config.Colours))
    i := 0
    for k, _ := range config.Colours {
        k_i, _ := strconv.ParseInt(k, 0, 64)
        mincount[i] = int(k_i)
	i++
    }

    sort.Sort(IntArray(mincount))

    data_in := map[string]int{ "HI": 1, "MN": 19, "MI": 20 }

    mapsvg_obj := svgxml.XML2SVG(mapsvg)
    if mapsvg_obj == nil {
        os.Exit(1)
    }

    re_fill, err := regexp.Compile(`(fill:#)......`)
    if err != nil {
        panic(err)
    }

    for state, state_count := range data_in {
	for _, mc := range mincount {
	    if state_count >= mc {
		element := svgxml.FindPathById(mapsvg_obj, state)
		if element != nil {
		    element.Style = string(re_fill.ReplaceAll([]byte(element.Style), []byte("${1}" + config.Colours[strconv.Itoa(mc)])))
		} else {
		    fmt.Fprintf(os.Stderr, "'%s' not found\n", state)
		}

	    }
	}
    }

    fmt.Println(string(svgxml.SVG2XML(mapsvg_obj, true)))

}

// vim:ai:sw=4:ts=8:
<|endoftext|>


<commit_before>package main

import (
	"strings"

	"github.com/prometheus/client_golang/prometheus"
)

func newMasterCollector(httpClient *httpClient) prometheus.Collector {
	metrics := map[prometheus.Collector]func(metricMap, prometheus.Collector) error{
		// CPU/Disk/Mem resources in free/used
		gauge("master", "cpus", "Current CPU resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/cpus_total"]
			used := m["master/cpus_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "cpus_revocable", "Current revocable CPU resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/cpus_revocable_total"]
			used := m["master/cpus_revocable_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "mem", "Current memory resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/mem_total"]
			used := m["master/mem_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "mem_revocable", "Current revocable memory resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/mem_revocable_total"]
			used := m["master/mem_revocable_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "disk", "Current disk resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/disk_total"]
			used := m["master/disk_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "disk_revocable", "Current disk resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/disk_revocable_total"]
			used := m["master/disk_revocable_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},

		// Master stats about uptime and election state
		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "elected",
			Help:      "1 if master is elected leader, 0 if not",
		}): func(m metricMap, c prometheus.Collector) error {
			elected, ok := m["master/elected"]
			if !ok {
				return notFoundInMap
			}
			c.(prometheus.Gauge).Set(elected)
			return nil
		},
		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "uptime_seconds",
			Help:      "Number of seconds the master process is running.",
		}): func(m metricMap, c prometheus.Collector) error {
			uptime, ok := m["master/uptime_secs"]
			if !ok {
				return notFoundInMap
			}
			c.(prometheus.Gauge).Set(uptime)
			return nil
		},
		// Master stats about agents
		counter("master", "slave_registration_events_total", "Total number of registration events on this master since it booted.", "event"): func(m metricMap, c prometheus.Collector) error {
			registrations := m["master/slave_registrations"]
			reregistrations := m["master/slave_reregistrations"]

			c.(*settableCounterVec).Set(registrations, "register")
			c.(*settableCounterVec).Set(reregistrations, "reregister")
			return nil
		},

		counter("master", "slave_removal_events_total", "Total number of removal events on this master since it booted.", "event"): func(m metricMap, c prometheus.Collector) error {
			scheduled := m["master/slave_shutdowns_scheduled"]
			canceled := m["master/slave_shutdowns_canceled"]
			completed := m["master/slave_shutdowns_completed"]
			removals := m["master/slave_removals"]

			c.(*settableCounterVec).Set(scheduled, "scheduled")
			c.(*settableCounterVec).Set(canceled, "canceled")
			c.(*settableCounterVec).Set(completed, "completed")
			c.(*settableCounterVec).Set(removals-completed, "died")
			return nil
		},
		gauge("master", "slaves_state", "Current number of slaves known to the master per connection and registration state.", "state"): func(m metricMap, c prometheus.Collector) error {
			active := m["master/slaves_active"]
			inactive := m["master/slaves_inactive"]
			disconnected := m["master/slaves_disconnected"]
			connected := m["master/slaves_connected"]

			c.(*prometheus.GaugeVec).WithLabelValues("active").Set(active)
			c.(*prometheus.GaugeVec).WithLabelValues("inactive").Set(inactive)
			c.(*prometheus.GaugeVec).WithLabelValues("disconnected").Set(disconnected)
			c.(*prometheus.GaugeVec).WithLabelValues("connected").Set(connected)
			return nil
		},

		// Master stats about frameworks
		gauge("master", "frameworks_state", "Current number of frames known to the master per connection and registration state.", "state"): func(m metricMap, c prometheus.Collector) error {
			active := m["master/frameworks_active"]
			connected := m["master/frameworks_connected"]
			inactive := m["master/frameworks_inactive"]
			disconnected := m["master/frameworks_disconnected"]

			c.(*prometheus.GaugeVec).WithLabelValues("active").Set(active)
			c.(*prometheus.GaugeVec).WithLabelValues("inactive").Set(inactive)
			c.(*prometheus.GaugeVec).WithLabelValues("disconnected").Set(disconnected)
			c.(*prometheus.GaugeVec).WithLabelValues("connected").Set(connected)
			return nil
		},
		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "offers_pending",
			Help:      "Current number of offers made by the master which aren't yet accepted or declined by frameworks.",
		}): func(m metricMap, c prometheus.Collector) error {
			offers := m["master/outstanding_offers"]

			c.(prometheus.Gauge).Set(offers)
			return nil
		},
		// Master stats about tasks
		counter("master", "task_states_exit_total", "Total number of tasks processed by exit state.", "state"): func(m metricMap, c prometheus.Collector) error {
			errored := m["master/tasks_error"]
			failed := m["master/tasks_failed"]
			finished := m["master/tasks_finished"]
			killed := m["master/tasks_killed"]
			lost := m["master/tasks_lost"]

			c.(*settableCounterVec).Set(errored, "errored")
			c.(*settableCounterVec).Set(failed, "failed")
			c.(*settableCounterVec).Set(finished, "finished")
			c.(*settableCounterVec).Set(killed, "killed")
			c.(*settableCounterVec).Set(lost, "lost")
			return nil
		},
		gauge("master", "task_states_current", "Current number of tasks by state.", "state"): func(m metricMap, c prometheus.Collector) error {
			running := m["master/tasks_running"]
			staging := m["master/tasks_staging"]
			starting := m["master/tasks_starting"]
			killing := m["master/tasks_killing"]
			unreachable := m["master/tasks_unreachable"]

			c.(*prometheus.GaugeVec).WithLabelValues("running").Set(running)
			c.(*prometheus.GaugeVec).WithLabelValues("staging").Set(staging)
			c.(*prometheus.GaugeVec).WithLabelValues("starting").Set(starting)
			c.(*prometheus.GaugeVec).WithLabelValues("killing").Set(killing)
			c.(*prometheus.GaugeVec).WithLabelValues("unreachable").Set(unreachable)
			return nil
		},

		// Master stats about messages
		counter("master", "messages_outcomes_total",
			"Total number of messages by outcome of operation and direction.",
			"source", "destination", "type", "outcome"): func(m metricMap, c prometheus.Collector) error {
			frameworkToExecutorValid := m["master/valid_framework_to_executor_messages"]
			frameworkToExecutorInvalid := m["master/invalid_framework_to_executor_messages"]
			executorToFrameworkValid := m["master/valid_executor_to_framework_messages"]
			executorToFrameworkInvalid := m["master/invalid_executor_to_framework_messages"]

			// status updates are sent from framework?(FIXME) to slave
			// status update acks are sent from slave to framework?
			statusUpdateAckValid := m["master/valid_status_update_acknowledgements"]
			statusUpdateAckInvalid := m["master/invalid_status_update_acknowledgements"]
			statusUpdateValid := m["master/valid_status_updates"]
			statusUpdateInvalid := m["master/invalid_status_updates"]

			c.(*settableCounterVec).Set(frameworkToExecutorValid, "framework", "executor", "", "valid")
			c.(*settableCounterVec).Set(frameworkToExecutorInvalid, "framework", "executor", "", "invalid")

			c.(*settableCounterVec).Set(executorToFrameworkValid, "executor", "framework", "", "valid")
			c.(*settableCounterVec).Set(executorToFrameworkInvalid, "executor", "framework", "", "invalid")

			// We consider a ack message simply as a message from slave to framework
			c.(*settableCounterVec).Set(statusUpdateValid, "framework", "slave", "status_update", "valid")
			c.(*settableCounterVec).Set(statusUpdateInvalid, "framework", "slave", "status_update", "invalid")
			c.(*settableCounterVec).Set(statusUpdateAckValid, "slave", "framework", "status_update", "valid")
			c.(*settableCounterVec).Set(statusUpdateAckInvalid, "slave", "framework", "status_update", "invalid")
			return nil
		},
		counter("master", "messages_type_total", "Total number of valid messages by type.", "type"): func(m metricMap, c prometheus.Collector) error {
			for k, v := range m {
				i := strings.Index("master/messages_", k)
				if i == -1 {
					continue
				}
				// FIXME: We expose things like messages_framework_to_executor twice
				c.(*settableCounterVec).Set(v, k[i:])
			}
			return nil
		},

		// Master stats about events
		gauge("master", "event_queue_length", "Current number of elements in event queue by type", "type"): func(m metricMap, c prometheus.Collector) error {
			dispatches := m["master/event_queue_dispatches"]
			httpRequests := m["master/event_queue_http_requests"]
			messages := m["master/event_queue_messages"]

			c.(*prometheus.GaugeVec).WithLabelValues("message").Set(messages)
			c.(*prometheus.GaugeVec).WithLabelValues("http_request").Set(httpRequests)
			c.(*prometheus.GaugeVec).WithLabelValues("dispatches").Set(dispatches)
			return nil
		},
	}
	return newMetricCollector(httpClient, metrics)
}
<commit_msg>Add mesos master allocation metrics.<commit_after>package main

import (
	"strings"

	"github.com/prometheus/client_golang/prometheus"
)

func newMasterCollector(httpClient *httpClient) prometheus.Collector {
	metrics := map[prometheus.Collector]func(metricMap, prometheus.Collector) error{
		// CPU/Disk/Mem resources in free/used
		gauge("master", "cpus", "Current CPU resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/cpus_total"]
			used := m["master/cpus_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "cpus_revocable", "Current revocable CPU resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/cpus_revocable_total"]
			used := m["master/cpus_revocable_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "mem", "Current memory resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/mem_total"]
			used := m["master/mem_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "mem_revocable", "Current revocable memory resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/mem_revocable_total"]
			used := m["master/mem_revocable_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "disk", "Current disk resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/disk_total"]
			used := m["master/disk_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},
		gauge("master", "disk_revocable", "Current disk resources in cluster.", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["master/disk_revocable_total"]
			used := m["master/disk_revocable_used"]

			c.(*prometheus.GaugeVec).WithLabelValues("free").Set(total - used)
			c.(*prometheus.GaugeVec).WithLabelValues("used").Set(used)
			return nil
		},

		// Master stats about uptime and election state
		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "elected",
			Help:      "1 if master is elected leader, 0 if not",
		}): func(m metricMap, c prometheus.Collector) error {
			elected, ok := m["master/elected"]
			if !ok {
				return notFoundInMap
			}
			c.(prometheus.Gauge).Set(elected)
			return nil
		},
		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "uptime_seconds",
			Help:      "Number of seconds the master process is running.",
		}): func(m metricMap, c prometheus.Collector) error {
			uptime, ok := m["master/uptime_secs"]
			if !ok {
				return notFoundInMap
			}
			c.(prometheus.Gauge).Set(uptime)
			return nil
		},
		// Master stats about agents
		counter("master", "slave_registration_events_total", "Total number of registration events on this master since it booted.", "event"): func(m metricMap, c prometheus.Collector) error {
			registrations := m["master/slave_registrations"]
			reregistrations := m["master/slave_reregistrations"]

			c.(*settableCounterVec).Set(registrations, "register")
			c.(*settableCounterVec).Set(reregistrations, "reregister")
			return nil
		},

		counter("master", "slave_removal_events_total", "Total number of removal events on this master since it booted.", "event"): func(m metricMap, c prometheus.Collector) error {
			scheduled := m["master/slave_shutdowns_scheduled"]
			canceled := m["master/slave_shutdowns_canceled"]
			completed := m["master/slave_shutdowns_completed"]
			removals := m["master/slave_removals"]

			c.(*settableCounterVec).Set(scheduled, "scheduled")
			c.(*settableCounterVec).Set(canceled, "canceled")
			c.(*settableCounterVec).Set(completed, "completed")
			c.(*settableCounterVec).Set(removals-completed, "died")
			return nil
		},
		gauge("master", "slaves_state", "Current number of slaves known to the master per connection and registration state.", "state"): func(m metricMap, c prometheus.Collector) error {
			active := m["master/slaves_active"]
			inactive := m["master/slaves_inactive"]
			disconnected := m["master/slaves_disconnected"]
			connected := m["master/slaves_connected"]

			c.(*prometheus.GaugeVec).WithLabelValues("active").Set(active)
			c.(*prometheus.GaugeVec).WithLabelValues("inactive").Set(inactive)
			c.(*prometheus.GaugeVec).WithLabelValues("disconnected").Set(disconnected)
			c.(*prometheus.GaugeVec).WithLabelValues("connected").Set(connected)
			return nil
		},

		// Master stats about frameworks
		gauge("master", "frameworks_state", "Current number of frames known to the master per connection and registration state.", "state"): func(m metricMap, c prometheus.Collector) error {
			active := m["master/frameworks_active"]
			connected := m["master/frameworks_connected"]
			inactive := m["master/frameworks_inactive"]
			disconnected := m["master/frameworks_disconnected"]

			c.(*prometheus.GaugeVec).WithLabelValues("active").Set(active)
			c.(*prometheus.GaugeVec).WithLabelValues("inactive").Set(inactive)
			c.(*prometheus.GaugeVec).WithLabelValues("disconnected").Set(disconnected)
			c.(*prometheus.GaugeVec).WithLabelValues("connected").Set(connected)
			return nil
		},
		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "offers_pending",
			Help:      "Current number of offers made by the master which aren't yet accepted or declined by frameworks.",
		}): func(m metricMap, c prometheus.Collector) error {
			offers := m["master/outstanding_offers"]

			c.(prometheus.Gauge).Set(offers)
			return nil
		},
		// Master stats about tasks
		counter("master", "task_states_exit_total", "Total number of tasks processed by exit state.", "state"): func(m metricMap, c prometheus.Collector) error {
			errored := m["master/tasks_error"]
			failed := m["master/tasks_failed"]
			finished := m["master/tasks_finished"]
			killed := m["master/tasks_killed"]
			lost := m["master/tasks_lost"]

			c.(*settableCounterVec).Set(errored, "errored")
			c.(*settableCounterVec).Set(failed, "failed")
			c.(*settableCounterVec).Set(finished, "finished")
			c.(*settableCounterVec).Set(killed, "killed")
			c.(*settableCounterVec).Set(lost, "lost")
			return nil
		},
		gauge("master", "task_states_current", "Current number of tasks by state.", "state"): func(m metricMap, c prometheus.Collector) error {
			running := m["master/tasks_running"]
			staging := m["master/tasks_staging"]
			starting := m["master/tasks_starting"]
			killing := m["master/tasks_killing"]
			unreachable := m["master/tasks_unreachable"]

			c.(*prometheus.GaugeVec).WithLabelValues("running").Set(running)
			c.(*prometheus.GaugeVec).WithLabelValues("staging").Set(staging)
			c.(*prometheus.GaugeVec).WithLabelValues("starting").Set(starting)
			c.(*prometheus.GaugeVec).WithLabelValues("killing").Set(killing)
			c.(*prometheus.GaugeVec).WithLabelValues("unreachable").Set(unreachable)
			return nil
		},

		// Master stats about messages
		counter("master", "messages_outcomes_total",
			"Total number of messages by outcome of operation and direction.",
			"source", "destination", "type", "outcome"): func(m metricMap, c prometheus.Collector) error {
			frameworkToExecutorValid := m["master/valid_framework_to_executor_messages"]
			frameworkToExecutorInvalid := m["master/invalid_framework_to_executor_messages"]
			executorToFrameworkValid := m["master/valid_executor_to_framework_messages"]
			executorToFrameworkInvalid := m["master/invalid_executor_to_framework_messages"]

			// status updates are sent from framework?(FIXME) to slave
			// status update acks are sent from slave to framework?
			statusUpdateAckValid := m["master/valid_status_update_acknowledgements"]
			statusUpdateAckInvalid := m["master/invalid_status_update_acknowledgements"]
			statusUpdateValid := m["master/valid_status_updates"]
			statusUpdateInvalid := m["master/invalid_status_updates"]

			c.(*settableCounterVec).Set(frameworkToExecutorValid, "framework", "executor", "", "valid")
			c.(*settableCounterVec).Set(frameworkToExecutorInvalid, "framework", "executor", "", "invalid")

			c.(*settableCounterVec).Set(executorToFrameworkValid, "executor", "framework", "", "valid")
			c.(*settableCounterVec).Set(executorToFrameworkInvalid, "executor", "framework", "", "invalid")

			// We consider a ack message simply as a message from slave to framework
			c.(*settableCounterVec).Set(statusUpdateValid, "framework", "slave", "status_update", "valid")
			c.(*settableCounterVec).Set(statusUpdateInvalid, "framework", "slave", "status_update", "invalid")
			c.(*settableCounterVec).Set(statusUpdateAckValid, "slave", "framework", "status_update", "valid")
			c.(*settableCounterVec).Set(statusUpdateAckInvalid, "slave", "framework", "status_update", "invalid")
			return nil
		},
		counter("master", "messages_type_total", "Total number of valid messages by type.", "type"): func(m metricMap, c prometheus.Collector) error {
			for k, v := range m {
				i := strings.Index("master/messages_", k)
				if i == -1 {
					continue
				}
				// FIXME: We expose things like messages_framework_to_executor twice
				c.(*settableCounterVec).Set(v, k[i:])
			}
			return nil
		},

		// Master stats about events
		gauge("master", "event_queue_length", "Current number of elements in event queue by type", "type"): func(m metricMap, c prometheus.Collector) error {
			dispatches := m["master/event_queue_dispatches"]
			httpRequests := m["master/event_queue_http_requests"]
			messages := m["master/event_queue_messages"]

			c.(*prometheus.GaugeVec).WithLabelValues("message").Set(messages)
			c.(*prometheus.GaugeVec).WithLabelValues("http_request").Set(httpRequests)
			c.(*prometheus.GaugeVec).WithLabelValues("dispatches").Set(dispatches)
			return nil
		},

		// Master stats about allocations

		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "allocation_run_ms_count",
			Help:      "Number of allocation algorithm time measurements in the window",
		}): func(m metricMap, c prometheus.Collector) error {
			count := m["allocator/mesos/allocation_runs"]
			c.(prometheus.Gauge).Set(count)
			return nil
		},

		gauge("master", "allocation_run_ms", "Time spent in allocation algorithm in ms.", "type"): func(m metricMap, c prometheus.Collector) error {
			mean := m["allocator/mesos/allocation_run_ms"]
			min := m["allocator/mesos/allocation_run_ms/min"]
			max := m["allocator/mesos/allocation_run_ms/max"]
			p50 := m["allocator/mesos/allocation_run_ms/p50"]
			p90 := m["allocator/mesos/allocation_run_ms/p90"]
			p95 := m["allocator/mesos/allocation_run_ms/p95"]
			p99 := m["allocator/mesos/allocation_run_ms/p99"]
			p999 := m["allocator/mesos/allocation_run_ms/p999"]
			p9999 := m["allocator/mesos/allocation_run_ms/p9999"]

			c.(*prometheus.GaugeVec).WithLabelValues("mean").Set(mean)
			c.(*prometheus.GaugeVec).WithLabelValues("min").Set(min)
			c.(*prometheus.GaugeVec).WithLabelValues("max").Set(max)
			c.(*prometheus.GaugeVec).WithLabelValues("p50").Set(p50)
			c.(*prometheus.GaugeVec).WithLabelValues("p90").Set(p90)
			c.(*prometheus.GaugeVec).WithLabelValues("p95").Set(p95)
			c.(*prometheus.GaugeVec).WithLabelValues("p99").Set(p99)
			c.(*prometheus.GaugeVec).WithLabelValues("p999").Set(p999)
			c.(*prometheus.GaugeVec).WithLabelValues("p9999").Set(p9999)
			return nil
		},

		prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "allocation_runs",
			Help:      "Number of times the allocation algorithm has run",
		}): func(m metricMap, c prometheus.Collector) error {
			count := m["allocator/mesos/allocation_runs"]
			c.(prometheus.Counter).Add(count)
			return nil
		},

		prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "allocation_run_latency_ms_count",
			Help:      "Number of allocation batch latency measurements in the window",
		}): func(m metricMap, c prometheus.Collector) error {
			count := m["allocator/mesos/allocation_runs"]
			c.(prometheus.Counter).Add(count)
			return nil
		},

		gauge("master", "allocation_run_latency_ms", "Allocation batch latency in ms.", "type"): func(m metricMap, c prometheus.Collector) error {
			mean := m["allocator/mesos/allocation_run_latency_ms"]
			min := m["allocator/mesos/allocation_run_latency_ms/min"]
			max := m["allocator/mesos/allocation_run_latency_ms/max"]
			p50 := m["allocator/mesos/allocation_run_latency_ms/p50"]
			p90 := m["allocator/mesos/allocation_run_latency_ms/p90"]
			p95 := m["allocator/mesos/allocation_run_latency_ms/p95"]
			p99 := m["allocator/mesos/allocation_run_latency_ms/p99"]
			p999 := m["allocator/mesos/allocation_run_latency_ms/p999"]
			p9999 := m["allocator/mesos/allocation_run_latency_ms/p9999"]

			c.(*prometheus.GaugeVec).WithLabelValues("mean").Set(mean)
			c.(*prometheus.GaugeVec).WithLabelValues("min").Set(min)
			c.(*prometheus.GaugeVec).WithLabelValues("max").Set(max)
			c.(*prometheus.GaugeVec).WithLabelValues("p50").Set(p50)
			c.(*prometheus.GaugeVec).WithLabelValues("p90").Set(p90)
			c.(*prometheus.GaugeVec).WithLabelValues("p95").Set(p95)
			c.(*prometheus.GaugeVec).WithLabelValues("p99").Set(p99)
			c.(*prometheus.GaugeVec).WithLabelValues("p999").Set(p999)
			c.(*prometheus.GaugeVec).WithLabelValues("p9999").Set(p9999)
			return nil
		},

		prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "mesos",
			Subsystem: "master",
			Name:      "event_queue_dispatches",
			Help:      "Number of dispatch events in the event queue.",
		}): func(m metricMap, c prometheus.Collector) error {
			count := m["allocator/mesos/event_queue_dispatches"]
			c.(prometheus.Gauge).Set(count)
			return nil
		},

		gauge("master", "allocator_resources_cpus", "Number of CPUs offered or allocated", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["allocator/mesos/resources/cpus/total"]
			offeredOrAllocated := m["allocator/mesos/resources/cpus/offered_or_allocated"]

			c.(*prometheus.GaugeVec).WithLabelValues("total").Set(total)
			c.(*prometheus.GaugeVec).WithLabelValues("offered_or_allocated").Set(offeredOrAllocated)
			return nil
		},

		gauge("master", "allocator_resources_disk", "Allocated or offered disk space in MB", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["allocator/mesos/resources/disk/total"]
			offeredOrAllocated := m["allocator/mesos/resources/disk/offered_or_allocated"]

			c.(*prometheus.GaugeVec).WithLabelValues("total").Set(total)
			c.(*prometheus.GaugeVec).WithLabelValues("offered_or_allocated").Set(offeredOrAllocated)
			return nil
		},

		gauge("master", "allocator_resources_mem", "Allocated or offered memory in MB", "type"): func(m metricMap, c prometheus.Collector) error {
			total := m["allocator/mesos/resources/mem/total"]
			offeredOrAllocated := m["allocator/mesos/resources/mem/offered_or_allocated"]

			c.(*prometheus.GaugeVec).WithLabelValues("total").Set(total)
			c.(*prometheus.GaugeVec).WithLabelValues("offered_or_allocated").Set(offeredOrAllocated)
			return nil
		},
	}
	return newMetricCollector(httpClient, metrics)
}
<|endoftext|>


<commit_before>package gomemcached

import (
	"encoding/binary"
	"fmt"
	"io"
)

// A memcached response
type MCResponse struct {
	// The command opcode of the command that sent the request
	Opcode CommandCode
	// The status of the response
	Status Status
	// The opaque sent in the request
	Opaque uint32
	// The CAS identifier (if applicable)
	Cas uint64
	// Extras, key, and body for this response
	Extras, Key, Body []byte
	// If true, this represents a fatal condition and we should hang up
	Fatal bool
}

// A debugging string representation of this response
func (res MCResponse) String() string {
	return fmt.Sprintf("{MCResponse status=%v keylen=%d, extralen=%d, bodylen=%d}",
		res.Status, len(res.Key), len(res.Extras), len(res.Body))
}

// Response as an error.
func (res MCResponse) Error() string {
	return fmt.Sprintf("MCResponse status=%v, opcode=%v, opaque=%v, msg: %s",
		res.Status, res.Opcode, res.Opaque, string(res.Body))
}

func errStatus(e error) Status {
	status := Status(0xffff)
	if res, ok := e.(MCResponse); ok {
		status = res.Status
	}
	if res, ok := e.(*MCResponse); ok {
		status = res.Status
	}
	return status
}

// True if this error represents a "not found" response.
func IsNotFound(e error) bool {
	return errStatus(e) == KEY_ENOENT
}

// False if this error isn't believed to be fatal to a connection.
func IsFatal(e error) bool {
	if e == nil {
		return false
	}
	switch errStatus(e) {
	case KEY_ENOENT, KEY_EEXISTS, EINVAL, NOT_STORED, TMPFAIL:
		return false
	}
	return true
}

// Number of bytes this response consumes on the wire.
func (res *MCResponse) Size() int {
	return HDR_LEN + len(res.Extras) + len(res.Key) + len(res.Body)
}

func (res *MCResponse) fillHeaderBytes(data []byte) int {
	pos := 0
	data[pos] = RES_MAGIC
	pos++
	data[pos] = byte(res.Opcode)
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2],
		uint16(len(res.Key)))
	pos += 2

	// 4
	data[pos] = byte(len(res.Extras))
	pos++
	data[pos] = 0
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2], uint16(res.Status))
	pos += 2

	// 8
	binary.BigEndian.PutUint32(data[pos:pos+4],
		uint32(len(res.Body)+len(res.Key)+len(res.Extras)))
	pos += 4

	// 12
	binary.BigEndian.PutUint32(data[pos:pos+4], res.Opaque)
	pos += 4

	// 16
	binary.BigEndian.PutUint64(data[pos:pos+8], res.Cas)
	pos += 8

	if len(res.Extras) > 0 {
		copy(data[pos:pos+len(res.Extras)], res.Extras)
		pos += len(res.Extras)
	}

	if len(res.Key) > 0 {
		copy(data[pos:pos+len(res.Key)], res.Key)
		pos += len(res.Key)
	}

	return pos
}

// Get just the header bytes for this response.
func (res *MCResponse) HeaderBytes() []byte {
	data := make([]byte, HDR_LEN+len(res.Extras)+len(res.Key))

	res.fillHeaderBytes(data)

	return data
}

// The actual bytes transmitted for this response.
func (res *MCResponse) Bytes() []byte {
	data := make([]byte, res.Size())

	pos := res.fillHeaderBytes(data)

	copy(data[pos:pos+len(res.Body)], res.Body)

	return data
}

// Send this response message across a writer.
func (res *MCResponse) Transmit(w io.Writer) (err error) {
	if len(res.Body) < 128 {
		_, err = w.Write(res.Bytes())
	} else {
		_, err = w.Write(res.HeaderBytes())
		if err == nil {
			_, err = w.Write(res.Body)
		}
	}
	return
}

// Fill this MCResponse with the data from this reader.
func (req *MCResponse) Receive(r io.Reader, hdrBytes []byte) error {
	if len(hdrBytes) < HDR_LEN {
		hdrBytes = []byte{
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0}
	}
	_, err := io.ReadFull(r, hdrBytes)
	if err != nil {
		return err
	}

	if hdrBytes[0] != RES_MAGIC && hdrBytes[0] != REQ_MAGIC {
		return fmt.Errorf("Bad magic: 0x%02x", hdrBytes[0])
	}

	klen := int(binary.BigEndian.Uint16(hdrBytes[2:4]))
	elen := int(hdrBytes[4])

	req.Opcode = CommandCode(hdrBytes[1])
	req.Status = Status(binary.BigEndian.Uint16(hdrBytes[6:8]))
	req.Opaque = binary.BigEndian.Uint32(hdrBytes[12:16])
	req.Cas = binary.BigEndian.Uint64(hdrBytes[16:24])

	bodyLen := int(binary.BigEndian.Uint32(hdrBytes[8:12])) - (klen + elen)

	buf := make([]byte, klen+elen+bodyLen)
	_, err = io.ReadFull(r, buf)
	if err == nil {
		req.Extras = buf[0:elen]
		req.Key = buf[elen : klen+elen]
		req.Body = buf[klen+elen:]
	}

	return err
}
<commit_msg>Consider EINVAL fatal<commit_after>package gomemcached

import (
	"encoding/binary"
	"fmt"
	"io"
)

// A memcached response
type MCResponse struct {
	// The command opcode of the command that sent the request
	Opcode CommandCode
	// The status of the response
	Status Status
	// The opaque sent in the request
	Opaque uint32
	// The CAS identifier (if applicable)
	Cas uint64
	// Extras, key, and body for this response
	Extras, Key, Body []byte
	// If true, this represents a fatal condition and we should hang up
	Fatal bool
}

// A debugging string representation of this response
func (res MCResponse) String() string {
	return fmt.Sprintf("{MCResponse status=%v keylen=%d, extralen=%d, bodylen=%d}",
		res.Status, len(res.Key), len(res.Extras), len(res.Body))
}

// Response as an error.
func (res MCResponse) Error() string {
	return fmt.Sprintf("MCResponse status=%v, opcode=%v, opaque=%v, msg: %s",
		res.Status, res.Opcode, res.Opaque, string(res.Body))
}

func errStatus(e error) Status {
	status := Status(0xffff)
	if res, ok := e.(MCResponse); ok {
		status = res.Status
	}
	if res, ok := e.(*MCResponse); ok {
		status = res.Status
	}
	return status
}

// True if this error represents a "not found" response.
func IsNotFound(e error) bool {
	return errStatus(e) == KEY_ENOENT
}

// False if this error isn't believed to be fatal to a connection.
func IsFatal(e error) bool {
	if e == nil {
		return false
	}
	switch errStatus(e) {
	case KEY_ENOENT, KEY_EEXISTS, NOT_STORED, TMPFAIL:
		return false
	}
	return true
}

// Number of bytes this response consumes on the wire.
func (res *MCResponse) Size() int {
	return HDR_LEN + len(res.Extras) + len(res.Key) + len(res.Body)
}

func (res *MCResponse) fillHeaderBytes(data []byte) int {
	pos := 0
	data[pos] = RES_MAGIC
	pos++
	data[pos] = byte(res.Opcode)
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2],
		uint16(len(res.Key)))
	pos += 2

	// 4
	data[pos] = byte(len(res.Extras))
	pos++
	data[pos] = 0
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2], uint16(res.Status))
	pos += 2

	// 8
	binary.BigEndian.PutUint32(data[pos:pos+4],
		uint32(len(res.Body)+len(res.Key)+len(res.Extras)))
	pos += 4

	// 12
	binary.BigEndian.PutUint32(data[pos:pos+4], res.Opaque)
	pos += 4

	// 16
	binary.BigEndian.PutUint64(data[pos:pos+8], res.Cas)
	pos += 8

	if len(res.Extras) > 0 {
		copy(data[pos:pos+len(res.Extras)], res.Extras)
		pos += len(res.Extras)
	}

	if len(res.Key) > 0 {
		copy(data[pos:pos+len(res.Key)], res.Key)
		pos += len(res.Key)
	}

	return pos
}

// Get just the header bytes for this response.
func (res *MCResponse) HeaderBytes() []byte {
	data := make([]byte, HDR_LEN+len(res.Extras)+len(res.Key))

	res.fillHeaderBytes(data)

	return data
}

// The actual bytes transmitted for this response.
func (res *MCResponse) Bytes() []byte {
	data := make([]byte, res.Size())

	pos := res.fillHeaderBytes(data)

	copy(data[pos:pos+len(res.Body)], res.Body)

	return data
}

// Send this response message across a writer.
func (res *MCResponse) Transmit(w io.Writer) (err error) {
	if len(res.Body) < 128 {
		_, err = w.Write(res.Bytes())
	} else {
		_, err = w.Write(res.HeaderBytes())
		if err == nil {
			_, err = w.Write(res.Body)
		}
	}
	return
}

// Fill this MCResponse with the data from this reader.
func (req *MCResponse) Receive(r io.Reader, hdrBytes []byte) error {
	if len(hdrBytes) < HDR_LEN {
		hdrBytes = []byte{
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0}
	}
	_, err := io.ReadFull(r, hdrBytes)
	if err != nil {
		return err
	}

	if hdrBytes[0] != RES_MAGIC && hdrBytes[0] != REQ_MAGIC {
		return fmt.Errorf("Bad magic: 0x%02x", hdrBytes[0])
	}

	klen := int(binary.BigEndian.Uint16(hdrBytes[2:4]))
	elen := int(hdrBytes[4])

	req.Opcode = CommandCode(hdrBytes[1])
	req.Status = Status(binary.BigEndian.Uint16(hdrBytes[6:8]))
	req.Opaque = binary.BigEndian.Uint32(hdrBytes[12:16])
	req.Cas = binary.BigEndian.Uint64(hdrBytes[16:24])

	bodyLen := int(binary.BigEndian.Uint32(hdrBytes[8:12])) - (klen + elen)

	buf := make([]byte, klen+elen+bodyLen)
	_, err = io.ReadFull(r, buf)
	if err == nil {
		req.Extras = buf[0:elen]
		req.Key = buf[elen : klen+elen]
		req.Body = buf[klen+elen:]
	}

	return err
}
<|endoftext|>


<commit_before>package gomemcached

import (
	"encoding/binary"
	"fmt"
	"io"
	"sync"
)

// MCResponse is memcached response
type MCResponse struct {
	// The command opcode of the command that sent the request
	Opcode CommandCode
	// The status of the response
	Status Status
	// The opaque sent in the request
	Opaque uint32
	// The CAS identifier (if applicable)
	Cas uint64
	// Extras, key, and body for this response
	Extras, FlexibleExtras, Key, Body []byte
	// If true, this represents a fatal condition and we should hang up
	Fatal bool
	// Datatype identifier
	DataType uint8
}

// A debugging string representation of this response
func (res MCResponse) String() string {
	return fmt.Sprintf("{MCResponse status=%v keylen=%d, extralen=%d, bodylen=%d, flexible=%v}",
		res.Status, len(res.Key), len(res.Extras), len(res.Body), res.FlexibleExtras)
}

// Response as an error.
func (res *MCResponse) Error() string {
	return fmt.Sprintf("MCResponse status=%v, opcode=%v, opaque=%v, msg: %s",
		res.Status, res.Opcode, res.Opaque, string(res.Body))
}

func errStatus(e error) Status {
	status := UNKNOWN_STATUS
	if res, ok := e.(*MCResponse); ok {
		status = res.Status
	}
	return status
}

// IsNotFound is true if this error represents a "not found" response.
func IsNotFound(e error) bool {
	return errStatus(e) == KEY_ENOENT
}

// IsFatal is false if this error isn't believed to be fatal to a connection.
func IsFatal(e error) bool {
	if e == nil {
		return false
	}
	_, ok := isFatal[errStatus(e)]
	if ok {
		return true
	}
	return false
}

// Size is number of bytes this response consumes on the wire.
func (res *MCResponse) Size() int {
	return HDR_LEN + len(res.Extras) + len(res.Key) + len(res.Body)
}

func (res *MCResponse) fillHeaderBytes(data []byte) int {
	pos := 0
	data[pos] = RES_MAGIC
	pos++
	data[pos] = byte(res.Opcode)
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2],
		uint16(len(res.Key)))
	pos += 2

	// 4
	data[pos] = byte(len(res.Extras))
	pos++
	// Data type
	if res.DataType != 0 {
		data[pos] = byte(res.DataType)
	} else {
		data[pos] = 0
	}
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2], uint16(res.Status))
	pos += 2

	// 8
	binary.BigEndian.PutUint32(data[pos:pos+4],
		uint32(len(res.Body)+len(res.Key)+len(res.Extras)))
	pos += 4

	// 12
	binary.BigEndian.PutUint32(data[pos:pos+4], res.Opaque)
	pos += 4

	// 16
	binary.BigEndian.PutUint64(data[pos:pos+8], res.Cas)
	pos += 8

	if len(res.Extras) > 0 {
		copy(data[pos:pos+len(res.Extras)], res.Extras)
		pos += len(res.Extras)
	}

	if len(res.Key) > 0 {
		copy(data[pos:pos+len(res.Key)], res.Key)
		pos += len(res.Key)
	}

	return pos
}

// HeaderBytes will get just the header bytes for this response.
func (res *MCResponse) HeaderBytes() []byte {
	data := make([]byte, HDR_LEN+len(res.Extras)+len(res.Key))

	res.fillHeaderBytes(data)

	return data
}

// Bytes will return the actual bytes transmitted for this response.
func (res *MCResponse) Bytes() []byte {
	data := make([]byte, res.Size())

	pos := res.fillHeaderBytes(data)

	copy(data[pos:pos+len(res.Body)], res.Body)

	return data
}

// Transmit will send this response message across a writer.
func (res *MCResponse) Transmit(w io.Writer) (n int, err error) {
	if len(res.Body) < 128 {
		n, err = w.Write(res.Bytes())
	} else {
		n, err = w.Write(res.HeaderBytes())
		if err == nil {
			m := 0
			m, err = w.Write(res.Body)
			m += n
		}
	}
	return
}

// Receive will fill this MCResponse with the data from this reader.
func (res *MCResponse) Receive(r io.Reader, hdrBytes []byte) (n int, err error) {
	return res.ReceiveWithBuf(r, hdrBytes, nil)
}

// ReceiveWithBuf takes an optional pre-allocated []byte buf which
// will be used if its capacity is large enough, otherwise a new
// []byte slice is allocated.
func (res *MCResponse) ReceiveWithBuf(r io.Reader, hdrBytes, buf []byte) (n int, err error) {
	if len(hdrBytes) < HDR_LEN {
		hdrBytes = []byte{
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0}
	}
	n, err = io.ReadFull(r, hdrBytes)
	if err != nil {
		return n, err
	}

	var klen, flen int
	switch hdrBytes[0] {
	case RES_MAGIC:
		fallthrough
	case REQ_MAGIC:
		klen = int(binary.BigEndian.Uint16(hdrBytes[2:4]))
	case FLEX_RES_MAGIC:
		flen = int(hdrBytes[2])
		klen = int(hdrBytes[3])
	default:
		return n, fmt.Errorf("bad magic: 0x%02x", hdrBytes[0])
	}
	elen := int(hdrBytes[4])

	res.Opcode = CommandCode(hdrBytes[1])
	res.DataType = uint8(hdrBytes[5])
	res.Status = Status(binary.BigEndian.Uint16(hdrBytes[6:8]))
	res.Opaque = binary.BigEndian.Uint32(hdrBytes[12:16])
	res.Cas = binary.BigEndian.Uint64(hdrBytes[16:24])

	bodyLen := int(binary.BigEndian.Uint32(hdrBytes[8:12])) - (klen + elen + flen)

	//defer function to debug the panic seen with MB-15557
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf(`Panic in Receive. Response %v \n
                        key len %v extra len %v bodylen %v`, res, klen, elen, bodyLen)
		}
	}()

	bufNeed := klen + elen + flen + bodyLen
	if buf != nil && cap(buf) >= bufNeed {
		buf = buf[0:bufNeed]
	} else {
		buf = make([]byte, bufNeed)
	}

	m, err := io.ReadFull(r, buf)
	if err == nil {
		if flen > 0 {
			res.FlexibleExtras = buf[0:flen]
			buf = buf[flen:]
		}
		res.Extras = buf[0:elen]
		res.Key = buf[elen : klen+elen]
		res.Body = buf[klen+elen:]
	}

	return n + m, err
}

func (res *MCResponse) ComputeUnits() (ru uint64, wu uint64) {
	if res.FlexibleExtras == nil || len(res.FlexibleExtras) == 0 {
		return
	}
	for i := 0; i < len(res.FlexibleExtras); {
		// TODO check: this seems to be the opposite of the documentation?
		l := res.FlexibleExtras[i] & 0x0f
		switch res.FlexibleExtras[i] >> 4 {
		case ComputeUnitsRead:
			ru = uint64(binary.BigEndian.Uint16(res.FlexibleExtras[i+1 : i+3]))
		case ComputeUnitsWrite:
			wu = uint64(binary.BigEndian.Uint16(res.FlexibleExtras[i+1 : i+3]))

		// ID escape: we need to skip the next byte
		case 15:
			i++
		}

		// data len is either 1..14, or 15 + next byte
		switch l {
		case 0:
			panic(fmt.Sprintf("Invalid Flexible Extras length received! %v", l))
		case 15:
			i = int(l + 1 + res.FlexibleExtras[i])
		default:
			i = int(l + 1)
		}
	}
	return
}

type MCResponsePool struct {
	pool *sync.Pool
}

func NewMCResponsePool() *MCResponsePool {
	rv := &MCResponsePool{
		pool: &sync.Pool{
			New: func() interface{} {
				return &MCResponse{}
			},
		},
	}

	return rv
}

func (this *MCResponsePool) Get() *MCResponse {
	return this.pool.Get().(*MCResponse)
}

func (this *MCResponsePool) Put(r *MCResponse) {
	if r == nil {
		return
	}

	r.Extras = nil
	r.Key = nil
	r.Body = nil
	r.Fatal = false

	this.pool.Put(r)
}

type StringMCResponsePool struct {
	pool *sync.Pool
	size int
}

func NewStringMCResponsePool(size int) *StringMCResponsePool {
	rv := &StringMCResponsePool{
		pool: &sync.Pool{
			New: func() interface{} {
				return make(map[string]*MCResponse, size)
			},
		},
		size: size,
	}

	return rv
}

func (this *StringMCResponsePool) Get() map[string]*MCResponse {
	return this.pool.Get().(map[string]*MCResponse)
}

func (this *StringMCResponsePool) Put(m map[string]*MCResponse) {
	if m == nil || len(m) > 2*this.size {
		return
	}

	for k := range m {
		m[k] = nil
		delete(m, k)
	}

	this.pool.Put(m)
}
<commit_msg>MB-53807 fix typo, flexible extra parsing does not advance properly<commit_after>package gomemcached

import (
	"encoding/binary"
	"fmt"
	"io"
	"sync"
)

// MCResponse is memcached response
type MCResponse struct {
	// The command opcode of the command that sent the request
	Opcode CommandCode
	// The status of the response
	Status Status
	// The opaque sent in the request
	Opaque uint32
	// The CAS identifier (if applicable)
	Cas uint64
	// Extras, key, and body for this response
	Extras, FlexibleExtras, Key, Body []byte
	// If true, this represents a fatal condition and we should hang up
	Fatal bool
	// Datatype identifier
	DataType uint8
}

// A debugging string representation of this response
func (res MCResponse) String() string {
	return fmt.Sprintf("{MCResponse status=%v keylen=%d, extralen=%d, bodylen=%d, flexible=%v}",
		res.Status, len(res.Key), len(res.Extras), len(res.Body), res.FlexibleExtras)
}

// Response as an error.
func (res *MCResponse) Error() string {
	return fmt.Sprintf("MCResponse status=%v, opcode=%v, opaque=%v, msg: %s",
		res.Status, res.Opcode, res.Opaque, string(res.Body))
}

func errStatus(e error) Status {
	status := UNKNOWN_STATUS
	if res, ok := e.(*MCResponse); ok {
		status = res.Status
	}
	return status
}

// IsNotFound is true if this error represents a "not found" response.
func IsNotFound(e error) bool {
	return errStatus(e) == KEY_ENOENT
}

// IsFatal is false if this error isn't believed to be fatal to a connection.
func IsFatal(e error) bool {
	if e == nil {
		return false
	}
	_, ok := isFatal[errStatus(e)]
	if ok {
		return true
	}
	return false
}

// Size is number of bytes this response consumes on the wire.
func (res *MCResponse) Size() int {
	return HDR_LEN + len(res.Extras) + len(res.Key) + len(res.Body)
}

func (res *MCResponse) fillHeaderBytes(data []byte) int {
	pos := 0
	data[pos] = RES_MAGIC
	pos++
	data[pos] = byte(res.Opcode)
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2],
		uint16(len(res.Key)))
	pos += 2

	// 4
	data[pos] = byte(len(res.Extras))
	pos++
	// Data type
	if res.DataType != 0 {
		data[pos] = byte(res.DataType)
	} else {
		data[pos] = 0
	}
	pos++
	binary.BigEndian.PutUint16(data[pos:pos+2], uint16(res.Status))
	pos += 2

	// 8
	binary.BigEndian.PutUint32(data[pos:pos+4],
		uint32(len(res.Body)+len(res.Key)+len(res.Extras)))
	pos += 4

	// 12
	binary.BigEndian.PutUint32(data[pos:pos+4], res.Opaque)
	pos += 4

	// 16
	binary.BigEndian.PutUint64(data[pos:pos+8], res.Cas)
	pos += 8

	if len(res.Extras) > 0 {
		copy(data[pos:pos+len(res.Extras)], res.Extras)
		pos += len(res.Extras)
	}

	if len(res.Key) > 0 {
		copy(data[pos:pos+len(res.Key)], res.Key)
		pos += len(res.Key)
	}

	return pos
}

// HeaderBytes will get just the header bytes for this response.
func (res *MCResponse) HeaderBytes() []byte {
	data := make([]byte, HDR_LEN+len(res.Extras)+len(res.Key))

	res.fillHeaderBytes(data)

	return data
}

// Bytes will return the actual bytes transmitted for this response.
func (res *MCResponse) Bytes() []byte {
	data := make([]byte, res.Size())

	pos := res.fillHeaderBytes(data)

	copy(data[pos:pos+len(res.Body)], res.Body)

	return data
}

// Transmit will send this response message across a writer.
func (res *MCResponse) Transmit(w io.Writer) (n int, err error) {
	if len(res.Body) < 128 {
		n, err = w.Write(res.Bytes())
	} else {
		n, err = w.Write(res.HeaderBytes())
		if err == nil {
			m := 0
			m, err = w.Write(res.Body)
			m += n
		}
	}
	return
}

// Receive will fill this MCResponse with the data from this reader.
func (res *MCResponse) Receive(r io.Reader, hdrBytes []byte) (n int, err error) {
	return res.ReceiveWithBuf(r, hdrBytes, nil)
}

// ReceiveWithBuf takes an optional pre-allocated []byte buf which
// will be used if its capacity is large enough, otherwise a new
// []byte slice is allocated.
func (res *MCResponse) ReceiveWithBuf(r io.Reader, hdrBytes, buf []byte) (n int, err error) {
	if len(hdrBytes) < HDR_LEN {
		hdrBytes = []byte{
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0, 0, 0}
	}
	n, err = io.ReadFull(r, hdrBytes)
	if err != nil {
		return n, err
	}

	var klen, flen int
	switch hdrBytes[0] {
	case RES_MAGIC:
		fallthrough
	case REQ_MAGIC:
		klen = int(binary.BigEndian.Uint16(hdrBytes[2:4]))
	case FLEX_RES_MAGIC:
		flen = int(hdrBytes[2])
		klen = int(hdrBytes[3])
	default:
		return n, fmt.Errorf("bad magic: 0x%02x", hdrBytes[0])
	}
	elen := int(hdrBytes[4])

	res.Opcode = CommandCode(hdrBytes[1])
	res.DataType = uint8(hdrBytes[5])
	res.Status = Status(binary.BigEndian.Uint16(hdrBytes[6:8]))
	res.Opaque = binary.BigEndian.Uint32(hdrBytes[12:16])
	res.Cas = binary.BigEndian.Uint64(hdrBytes[16:24])

	bodyLen := int(binary.BigEndian.Uint32(hdrBytes[8:12])) - (klen + elen + flen)

	//defer function to debug the panic seen with MB-15557
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf(`Panic in Receive. Response %v \n
                        key len %v extra len %v bodylen %v`, res, klen, elen, bodyLen)
		}
	}()

	bufNeed := klen + elen + flen + bodyLen
	if buf != nil && cap(buf) >= bufNeed {
		buf = buf[0:bufNeed]
	} else {
		buf = make([]byte, bufNeed)
	}

	m, err := io.ReadFull(r, buf)
	if err == nil {
		if flen > 0 {
			res.FlexibleExtras = buf[0:flen]
			buf = buf[flen:]
		}
		res.Extras = buf[0:elen]
		res.Key = buf[elen : klen+elen]
		res.Body = buf[klen+elen:]
	}

	return n + m, err
}

func (res *MCResponse) ComputeUnits() (ru uint64, wu uint64) {
	if res.FlexibleExtras == nil || len(res.FlexibleExtras) == 0 {
		return
	}
	for i := 0; i < len(res.FlexibleExtras); {
		// TODO check: this seems to be the opposite of the documentation?
		l := res.FlexibleExtras[i] & 0x0f
		switch res.FlexibleExtras[i] >> 4 {
		case ComputeUnitsRead:
			ru = uint64(binary.BigEndian.Uint16(res.FlexibleExtras[i+1 : i+3]))
		case ComputeUnitsWrite:
			wu = uint64(binary.BigEndian.Uint16(res.FlexibleExtras[i+1 : i+3]))

		// ID escape: we need to skip the next byte
		case 15:
			i++
		}

		// data len is either 1..14, or 15 + next byte
		switch l {
		case 0:
			panic(fmt.Sprintf("Invalid Flexible Extras length received! %v", l))
		case 15:
			i = int(l + 1 + res.FlexibleExtras[i])
		default:
			i += int(l + 1)
		}
	}
	return
}

type MCResponsePool struct {
	pool *sync.Pool
}

func NewMCResponsePool() *MCResponsePool {
	rv := &MCResponsePool{
		pool: &sync.Pool{
			New: func() interface{} {
				return &MCResponse{}
			},
		},
	}

	return rv
}

func (this *MCResponsePool) Get() *MCResponse {
	return this.pool.Get().(*MCResponse)
}

func (this *MCResponsePool) Put(r *MCResponse) {
	if r == nil {
		return
	}

	r.Extras = nil
	r.Key = nil
	r.Body = nil
	r.Fatal = false

	this.pool.Put(r)
}

type StringMCResponsePool struct {
	pool *sync.Pool
	size int
}

func NewStringMCResponsePool(size int) *StringMCResponsePool {
	rv := &StringMCResponsePool{
		pool: &sync.Pool{
			New: func() interface{} {
				return make(map[string]*MCResponse, size)
			},
		},
		size: size,
	}

	return rv
}

func (this *StringMCResponsePool) Get() map[string]*MCResponse {
	return this.pool.Get().(map[string]*MCResponse)
}

func (this *StringMCResponsePool) Put(m map[string]*MCResponse) {
	if m == nil || len(m) > 2*this.size {
		return
	}

	for k := range m {
		m[k] = nil
		delete(m, k)
	}

	this.pool.Put(m)
}
<|endoftext|>


<commit_before>package mgorus

import (
	"fmt"

	"github.com/Sirupsen/logrus"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type hooker struct {
	c *mgo.Collection
}

type M bson.M

func NewHooker(mgoUrl, db, collection string) (*hooker, error) {
	session, err := mgo.Dial(mgoUrl)
	if err != nil {
		return nil, err
	}

	return &hooker{c: session.DB(db).C(collection)}, nil
}

func (h *hooker) Fire(entry *logrus.Entry) error {
	entry.Data["Level"] = entry.Level.String()
	entry.Data["Time"] = entry.Time
	entry.Data["Message"] = entry.Message
	if errData, ok := entry.Data[logrus.ErrorKey]; ok  {
		if err, ok := errData.(error); ok && entry.Data[logrus.ErrorKey]!=nil {
				entry.Data[logrus.ErrorKey] = err.Error()
		}
	}
	mgoErr := h.c.Insert(M(entry.Data))
	if mgoErr != nil {
		return fmt.Errorf("Failed to send log entry to mongodb: %s", mgoErr)
	}

	return nil
}

func (h *hooker) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}
<commit_msg>gofmt<commit_after>package mgorus

import (
	"fmt"

	"github.com/Sirupsen/logrus"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type hooker struct {
	c *mgo.Collection
}

type M bson.M

func NewHooker(mgoUrl, db, collection string) (*hooker, error) {
	session, err := mgo.Dial(mgoUrl)
	if err != nil {
		return nil, err
	}

	return &hooker{c: session.DB(db).C(collection)}, nil
}

func (h *hooker) Fire(entry *logrus.Entry) error {
	entry.Data["Level"] = entry.Level.String()
	entry.Data["Time"] = entry.Time
	entry.Data["Message"] = entry.Message
	if errData, ok := entry.Data[logrus.ErrorKey]; ok {
		if err, ok := errData.(error); ok && entry.Data[logrus.ErrorKey] != nil {
			entry.Data[logrus.ErrorKey] = err.Error()
		}
	}
	mgoErr := h.c.Insert(M(entry.Data))
	if mgoErr != nil {
		return fmt.Errorf("Failed to send log entry to mongodb: %s", mgoErr)
	}

	return nil
}

func (h *hooker) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}
<|endoftext|>


<commit_before>package minnow

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/binary"
	"errors"
	"hash"
	"io"
	"log"
	"sync"
	"time"

	cryptoRand "crypto/rand"
	mathRand "math/rand"
)

var ErrAlreadyClosed = errors.New("Writer has already been closed")
var ErrMaxLengthExceeded = errors.New("Maximum message length (1<<64 - 1) exceeded")

type Packet struct {
	Header  PacketHeader
	Payload []byte
}

type PacketHeader struct {
	SequenceN uint64
	Mac       [64]byte
	Size      uint64
}

type MessageWriteCloser struct {
	closed      bool
	hash        hash.Hash
	message     []byte
	destination io.WriteCloser
	wg          *sync.WaitGroup
}

func NewMessageWriteCloser(secret []byte, w io.WriteCloser) *MessageWriteCloser {
	return &MessageWriteCloser{
		closed:      false,
		hash:        hmac.New(sha512.New, secret),
		message:     make([]byte, 0),
		destination: w,
		wg:          new(sync.WaitGroup),
	}
}

// Writes a message fragment to the internal buffer
// Return value is always (len(m), nil)
func (mw *MessageWriteCloser) Write(m []byte) (int, error) {
	mw.message = append(mw.message, m...)
	return len(m), nil
}

// Close is when we actually get the end of the message and can write it
func (mw *MessageWriteCloser) Close() error {
	if !mw.closed {
		log.Printf("Writing message:\n%s", string(mw.message))
		mw.closed = true
		// Write chaff asynchronously
		mw.wg.Add(1)
		go binaryChaff(mw.destination, mw.wg.Done)

		// Write the real message
		_, err := mw.writeMessage()

		// Wait for chaff to finish writing before closing connection
		mw.wg.Wait()

		// Either close and handle error or just close
		if err != nil {
			// Always close the destination
			mw.destination.Close()
			return err
		}
		return mw.destination.Close()
	}
	return ErrAlreadyClosed
}

func (mw *MessageWriteCloser) writeMessage() (n int, err error) {
	log.Printf("Writing message <<%s>> len: %d cap: %d", string(mw.message), len(mw.message), cap(mw.message))
	for n, _ = range mw.message {
		_, err = mw.writePacket(uint64(n), mw.message[n:n+1])
		log.Printf("Wrote seq %d: %v with err: %v", n, mw.message[n:n+1], err)
		if err != nil {
			return
		}
		time.Sleep(time.Duration(mathRand.Intn(10)) * time.Millisecond)
	}
	return
}

func (mw *MessageWriteCloser) writePacket(seqn uint64, packet []byte) (int, error) {
	var mbuf [64]byte

	mw.hash.Write(packet)
	defer mw.hash.Reset()

	m := mw.hash.Sum(nil)

	for i, _ := range mbuf {
		mbuf[i] = m[i]
	}

	h := PacketHeader{
		SequenceN: seqn,
		Mac:       mbuf,
		Size:      uint64(len(packet)),
	}

	err := binary.Write(mw.destination, binary.BigEndian, h)
	if err != nil {
		return 0, err
	}

	n, err := mw.destination.Write(packet)
	if err != nil {
		return n, err
	}

	return len(packet), nil
}

type MessageReader struct {
	hash    hash.Hash
	message []byte
	reader  io.Reader
	closed  bool
}

func NewMessageReader(secret []byte, r io.Reader) *MessageReader {
	return &MessageReader{
		hash:    hmac.New(sha512.New, secret),
		message: make([]byte, 0),
		reader:  r,
		closed:  false,
	}
}

func (r *MessageReader) ReadAll() []byte {
	for {
		res := new(PacketHeader)
		err := binary.Read(r.reader, binary.BigEndian, res)
		if err != nil {
			if err != io.EOF {
				// Probably shouldn't Fatal on this...
				log.Fatalf("Received an error reading header: %w\n", err)
			}
			r.closed = true
			break
		}

		d := make([]byte, res.Size)
		_, err = io.ReadFull(r.reader, d)
		if err != nil {
			if err != io.EOF {
				log.Fatalf("Received error reading message body: %w\n", err)
			}
			r.closed = true
			break
		}
		if r.matches(d, res.Mac) {
			r.message = append(r.message, d...)
		}
	}
	return r.message
}

func (r *MessageReader) matches(data []byte, provided [64]byte) bool {
	r.hash.Write(data)
	defer r.hash.Reset()

	m := r.hash.Sum(nil)

	var mbuf [64]byte
	for i, _ := range mbuf {
		mbuf[i] = m[i]
	}

	return mbuf == provided
}

func binaryChaff(w io.Writer, done func()) {
	defer done()

	var err error
	ph := &PacketHeader{Size: 1}
	macbuf := make([]byte, 64, 64)
	payloadbuf := make([]byte, 1, 1)
	sz := uint64(0)
	for sz == 0 {
		//FIXME Need to base amount of chaff of message size without
		//      hinting at message size
		sz = uint64(mathRand.Int63()) / 1000000000000000
	}

	log.Printf("Writing %d chaff packets", sz)

	for ; ph.SequenceN < sz; ph.SequenceN++ {
		_, err = cryptoRand.Read(macbuf)
		if err != nil {
			return
		}

		for i, _ := range macbuf {
			ph.Mac[i] = macbuf[i]
		}

		err = binary.Write(w, binary.BigEndian, ph)
		if err != nil {
			log.Printf("Failed to write chaff header: %s", err)
			return
		}

		_, err = cryptoRand.Read(payloadbuf)
		if err != nil {
			return
		}

		_, err = w.Write(payloadbuf)
		if err != nil {
			return
		}
		time.Sleep(time.Duration(mathRand.Intn(2)) * time.Millisecond)
	}
	log.Printf("Wrote %d chaff packets", ph.SequenceN)

}
<commit_msg>Run packets through channels for random chaff distribution<commit_after>package minnow

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/binary"
	"errors"
	"hash"
	"io"
	"log"
	"sync"

	cryptoRand "crypto/rand"
	mathRand "math/rand"
)

var ErrAlreadyClosed = errors.New("Writer has already been closed")
var ErrMaxLengthExceeded = errors.New("Maximum message length (1<<64 - 1) exceeded")

type FakePacketResult struct {
    Hasher hash.Hash
    Message []byte
}

type Packet struct {
	Header  PacketHeader
	Payload []byte
}

type PacketHeader struct {
	SequenceN uint64
	Mac       [64]byte
	Size      uint64
}

type MessageWriteCloser struct {
	closed      bool
	hash        hash.Hash
	message     []byte
	destination io.WriteCloser
	wg          *sync.WaitGroup
}

func NewMessageWriteCloser(secret []byte, w io.WriteCloser) *MessageWriteCloser {
	return &MessageWriteCloser{
		closed:      false,
		hash:        hmac.New(sha512.New, secret),
		message:     make([]byte, 0),
		destination: w,
		wg:          new(sync.WaitGroup),
	}
}

// Writes a message fragment to the internal buffer
// Return value is always (len(m), nil)
func (mw *MessageWriteCloser) Write(m []byte) (int, error) {
	mw.message = append(mw.message, m...)
	return len(m), nil
}

// Close is when we actually get the end of the message and can write it
func (mw *MessageWriteCloser) Close() error {
	if !mw.closed {
        defer mw.destination.Close()
		log.Printf("Writing message:\n%s", string(mw.message))
		mw.closed = true
        packetchans := make([]chan *Packet, 11)
        packetchans[0] = sequentialPacketChannel(mw.message, mw.hash)
		// Write chaff asynchronously
        // FIXME: Random number of chaffers
        for i := 1; i < 11; i++ {
            // FIXME: Length should be randomized
            faker, err := fakeMessage(len(mw.message))
            if err != nil {
                return err
            }
            packetchans[i] = sequentialPacketChannel(faker.Message, faker.Hasher)
        }

        outpackets := randomPacketRouter(packetchans)

        // FIXME: Randomize the packets/order
        for p := range outpackets {
            err := mw.writePacket(p)
            if err != nil {
                return err
            }
        }

		return mw.destination.Close()
	}
	return ErrAlreadyClosed
}

func (mw *MessageWriteCloser) writePacket(packet *Packet) error {
    log.Printf("Writing packet: #%d - %+v\n", packet.Header.SequenceN, packet.Payload)
	err := binary.Write(mw.destination, binary.BigEndian, packet.Header)
	if err != nil {
		return err
	}

	_, err = mw.destination.Write(packet.Payload)
	if err != nil {
		return err
	}

	return nil
}

type MessageReader struct {
	hash    hash.Hash
	message []byte
	reader  io.Reader
	closed  bool
}

func NewMessageReader(secret []byte, r io.Reader) *MessageReader {
	return &MessageReader{
		hash:    hmac.New(sha512.New, secret),
		message: make([]byte, 0),
		reader:  r,
		closed:  false,
	}
}

func (r *MessageReader) ReadAll() []byte {
	for {
		res := new(PacketHeader)
		err := binary.Read(r.reader, binary.BigEndian, res)
		if err != nil {
			if err != io.EOF {
				// Probably shouldn't Fatal on this...
				log.Fatalf("Received an error reading header: %w\n", err)
			}
			r.closed = true
			break
		}

		d := make([]byte, res.Size)
		_, err = io.ReadFull(r.reader, d)
		if err != nil {
			if err != io.EOF {
				log.Fatalf("Received error reading message body: %w\n", err)
			}
			r.closed = true
			break
		}
		if r.matches(d, res.Mac) {
			r.message = append(r.message, d...)
		}
	}
	return r.message
}

func (r *MessageReader) matches(data []byte, provided [64]byte) bool {
	r.hash.Write(data)
	defer r.hash.Reset()

	m := r.hash.Sum(nil)

	var mbuf [64]byte
	for i, _ := range mbuf {
		mbuf[i] = m[i]
	}

	return mbuf == provided
}

// Send len(data) # of sequential packets, hashed with the provided hash.Hash
// down the provided channel
func sequentialPacketChannel(data []byte, hasher hash.Hash) chan *Packet {
    // FIXME: Packet size is hard-coded at 1
    outchan := make(chan *Packet)
    go func() {
        for i, _ := range data {
            var mbuf [64]byte

            hasher.Write(data[i:i+1])
            m := hasher.Sum(nil)
            for i, _ := range mbuf {
                mbuf[i] = m[i]
            }
            hasher.Reset()

            h := PacketHeader{
                SequenceN: uint64(i),
                Mac:       mbuf,
                Size:      uint64(1),
            }
            outchan<- &Packet{h, data[i:i+1]}
        }

        outchan<- nil
    }()

    return outchan
}

// Create a mock hasher and buffer
func fakeMessage(packetSize int) (*FakePacketResult, error) {
    var err error

    payloadbuf := make([]byte, packetSize, packetSize)

    sz := uint64(0)
    for sz == 0 {
        sz = uint64(mathRand.Int63()) / 1000000000000000
    }
    hashseed := make([]byte, sz, sz)

    _, err = cryptoRand.Read(hashseed)
    if err != nil {
        return nil, err
    }

    h := hmac.New(sha512.New, hashseed)

    _, err = cryptoRand.Read(payloadbuf)
    if err != nil {
        return nil, err
    }

    return &FakePacketResult{h, payloadbuf}, nil
}

// Radomly select packets from the provided channels until they
// are all depleted (received a nil)
func randomPacketRouter(inchans []chan *Packet) chan *Packet {
    outchan := make(chan *Packet)
    go func() {
        for len(inchans) > 0 {
            ndx := mathRand.Intn(len(inchans))
            val := <-inchans[ndx]
            // This packet provider is done
            if val == nil {
                inchans = append(inchans[:ndx], inchans[ndx+1:]...)
            } else {
                outchan<- val
            }
        }
        close(outchan)
    }()

    return outchan
}
<|endoftext|>


<commit_before><commit_msg>Bring back context cancelled TODO<commit_after><|endoftext|>


<commit_before>// Copyright 2012 commandmocker authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package commandmocker

import (
	"errors"
	"fmt"
	"os"
	"path"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"
)

var source = `#!/bin/sh -e

output="{{.}}"
echo -n "${output}"
touch $(dirname ${0})/.ran
`

var pathMutex sync.Mutex

// Add creates a temporary directory containing an executable file named "name"
// that prints "output" when executed. It also adds the temporary directory to
// the first position of $PATH.
//
// It returns the temporary directory path (for future removing, using the
// Remove function) and an error if any happen.
func Add(name, output string) (tempdir string, err error) {
	tempdir = path.Join(os.TempDir(), "commandmocker+" + time.Now().Format("20060102150405999999999"))
	_, err = os.Stat(tempdir)
	for !os.IsNotExist(err) {
		tempdir = path.Join(os.TempDir(), "commandmocker+" + time.Now().Format("20060102150405999999999"))
		_, err = os.Stat(tempdir)
	}
	err = os.MkdirAll(tempdir, 0777)
	if err != nil {
		return
	}
	f, err := os.OpenFile(path.Join(tempdir, name), syscall.O_WRONLY|syscall.O_CREAT|syscall.O_TRUNC, 0755)
	if err != nil {
		return
	}
	defer f.Close()
	t, err := template.New(name).Parse(source)
	if err != nil {
		return
	}
	err = t.Execute(f, output)
	if err != nil {
		return
	}
	pathMutex.Lock()
	defer pathMutex.Unlock()
	path := os.Getenv("PATH")
	path = tempdir + ":" + path
	err = os.Setenv("PATH", path)
	return
}

// Ran indicates whether the mocked executable was called or not.
//
// It just checks if the given tempdir contains a .ran file.
func Ran(tempdir string) bool {
	p := path.Join(tempdir, ".ran")
	_, err := os.Stat(p)
	return err == nil || !os.IsNotExist(err)
}

// Remove removes the tempdir from $PATH and from file system.
//
// This function is intended only to undo what Add does. It returns error if
// the given tempdir is not a temporary directory.
func Remove(tempdir string) error {
	if !strings.HasPrefix(tempdir, os.TempDir()) {
		return errors.New("Remove can only remove temporary directories, tryied to remove " + tempdir)
	}
	pathMutex.Lock()
	path := os.Getenv("PATH")
	index := strings.Index(path, tempdir)
	if index < 0 {
		pathMutex.Unlock()
		return errors.New(fmt.Sprintf("%s is not in $PATH", tempdir))
	}
	path = path[:index] + path[index+len(tempdir)+1:]
	err := os.Setenv("PATH", path)
	pathMutex.Unlock()
	if err != nil {
		return err
	}
	return os.RemoveAll(tempdir)
}
<commit_msg>mocker: using bash instead of sh<commit_after>// Copyright 2012 commandmocker authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package commandmocker

import (
	"errors"
	"fmt"
	"os"
	"path"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"
)

var source = `#!/bin/bash -e

output="{{.}}"
echo -n "${output}"
touch $(dirname ${0})/.ran
`

var pathMutex sync.Mutex

// Add creates a temporary directory containing an executable file named "name"
// that prints "output" when executed. It also adds the temporary directory to
// the first position of $PATH.
//
// It returns the temporary directory path (for future removing, using the
// Remove function) and an error if any happen.
func Add(name, output string) (tempdir string, err error) {
	tempdir = path.Join(os.TempDir(), "commandmocker+" + time.Now().Format("20060102150405999999999"))
	_, err = os.Stat(tempdir)
	for !os.IsNotExist(err) {
		tempdir = path.Join(os.TempDir(), "commandmocker+" + time.Now().Format("20060102150405999999999"))
		_, err = os.Stat(tempdir)
	}
	err = os.MkdirAll(tempdir, 0777)
	if err != nil {
		return
	}
	f, err := os.OpenFile(path.Join(tempdir, name), syscall.O_WRONLY|syscall.O_CREAT|syscall.O_TRUNC, 0755)
	if err != nil {
		return
	}
	defer f.Close()
	t, err := template.New(name).Parse(source)
	if err != nil {
		return
	}
	err = t.Execute(f, output)
	if err != nil {
		return
	}
	pathMutex.Lock()
	defer pathMutex.Unlock()
	path := os.Getenv("PATH")
	path = tempdir + ":" + path
	err = os.Setenv("PATH", path)
	return
}

// Ran indicates whether the mocked executable was called or not.
//
// It just checks if the given tempdir contains a .ran file.
func Ran(tempdir string) bool {
	p := path.Join(tempdir, ".ran")
	_, err := os.Stat(p)
	return err == nil || !os.IsNotExist(err)
}

// Remove removes the tempdir from $PATH and from file system.
//
// This function is intended only to undo what Add does. It returns error if
// the given tempdir is not a temporary directory.
func Remove(tempdir string) error {
	if !strings.HasPrefix(tempdir, os.TempDir()) {
		return errors.New("Remove can only remove temporary directories, tryied to remove " + tempdir)
	}
	pathMutex.Lock()
	path := os.Getenv("PATH")
	index := strings.Index(path, tempdir)
	if index < 0 {
		pathMutex.Unlock()
		return errors.New(fmt.Sprintf("%s is not in $PATH", tempdir))
	}
	path = path[:index] + path[index+len(tempdir)+1:]
	err := os.Setenv("PATH", path)
	pathMutex.Unlock()
	if err != nil {
		return err
	}
	return os.RemoveAll(tempdir)
}
<|endoftext|>


<commit_before>package posts

import (
	"encoding/json"
	"io/ioutil"
	"testing"
	"time"

	"github.com/boatilus/peppercorn/db"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	rethink "gopkg.in/dancannon/gorethink.v2"
)

const tableName = "posts_test"

type doc struct {
	Active  bool
	Author  string
	Content string
	Time    int64
}

var docs []doc // Stores test data read in from JSON

func makePostFromDoc(d doc) Post {
	return Post{
		Active:  d.Active,
		Author:  d.Author,
		Content: d.Content,
		Time:    time.Unix(d.Time, 0),
	}
}

func setupDB() {
	if !db.Session.IsConnected() {
		panic("No DB connected")
	}

	rethink.DBCreate("peppercorn").RunWrite(db.Session)

	peppercorn := rethink.DB(db.Name)

	c, err := peppercorn.TableList().Contains(tableName).Run(db.Session)
	if err != nil {
		panic(err)
	}

	var hasTable bool

	err = c.One(&hasTable)
	if err != nil {
		panic(err)
	}

	table := peppercorn.Table(tableName)

	if !hasTable {
		_, err := peppercorn.TableCreate(tableName).RunWrite(db.Session)
		if err != nil {
			panic(err)
		}

		table.IndexCreate("active").Run(db.Session)
		table.IndexCreate("author").Run(db.Session)
		table.IndexCreate("time").Run(db.Session)
		table.IndexWait().Run(db.Session)
	} else {
		// Due to a lack of mocking in gorethink, we'll tear down the test data and repopulate on each
		// run of the tests.
		table.Delete().RunWrite(db.Session)
	}

	bytes, err := ioutil.ReadFile("posts.test_data.json")

	if err := json.Unmarshal(bytes, &docs); err != nil {
		panic(err)
	}

	if len(docs) != 7 {
		panic(err)
	}

	posts := make([]Post, len(docs))

	for i := range docs {
		posts[i].Active = docs[i].Active
		posts[i].Author = docs[i].Author
		posts[i].Content = docs[i].Content
		posts[i].Time = time.Unix(docs[i].Time, 0)
	}

	if _, err := table.Insert(posts).RunWrite(db.Session); err != nil {
		panic(err)
	}

	cursor, err := table.Count().Run(db.Session)

	if err != nil {
		panic(err)
	}

	var n int

	cursor.One(&n)
	cursor.Close()

	if n != 7 {
		panic(err)
	}
}

func init() {
	viper.Set("db.posts_table", "posts_test")

	var err error

	if db.Session, err = rethink.Connect(rethink.ConnectOpts{Address: "localhost:28015"}); err != nil {
		panic(err)
	}

	setupDB()
}

///////////
// Tests //
///////////

func TestNew(t *testing.T) {
	assert := assert.New(t)

	type data struct {
		author  string
		content string
	}

	passCases := []data{
		{"user", "content"},
		{"_", "_"},
	}

	for _, c := range passCases {
		got, err := New(c.author, c.content)

		assert.Nil(err)
		assert.Empty(got.ID)
		assert.True(got.Active)
		assert.Equal(c.author, got.Author)
		assert.Equal(c.content, got.Content)
	}

	failCases := []data{
		{"", "content"},
		{"user", ""},
		{"", ""},
	}

	for _, c := range failCases {
		_, err := New(c.author, c.content)

		assert.NotNil(err)
	}
}

func TestCount(t *testing.T) {
	n, err := Count()

	assert.Nil(t, err)
	assert.Equal(t, n, 6)
}

func TestCountAll(t *testing.T) {
	n, err := CountAll()

	assert.Nil(t, err)
	assert.Equal(t, n, 7)
}

func TestGetRange(t *testing.T) {
	assert := assert.New(t)

	cases := []struct {
		first uint64
		limit uint64
		want  []Post
	}{
		{1, 2, []Post{makePostFromDoc(docs[0]), makePostFromDoc(docs[1])}},
		{3, 2, []Post{makePostFromDoc(docs[2]), makePostFromDoc(docs[3])}},
		{2, 3, []Post{makePostFromDoc(docs[1]), makePostFromDoc(docs[2]), makePostFromDoc(docs[3])}},
		// The 'first' argument is locked to 1 if < 1, so we should check that we get posts 1 and 2...
		{0, 2, []Post{makePostFromDoc(docs[0]), makePostFromDoc(docs[1])}},
	}

	for _, c := range cases {
		got, err := GetRange(c.first, c.limit)

		assert.Nil(err)

		for i := range got {
			g := got[i]
			w := c.want[i]

			assert.Equal(g.Active, w.Active)
			assert.Equal(g.Author, w.Author)
			assert.Equal(g.Content, w.Content)
			assert.True(g.Time.Equal(w.Time))
		}
	}
}

func TestGetOne(t *testing.T) {
	assert := assert.New(t)

	cases := []struct {
		in   uint64
		want Post
	}{
		{1, makePostFromDoc(docs[0])},
		{3, makePostFromDoc(docs[2])},
	}

	for _, c := range cases {
		got, err := GetOne(c.in)

		assert.Nil(err)

		assert.Equal(got.Active, c.want.Active)
		assert.Equal(got.Author, c.want.Author)
		assert.Equal(got.Content, c.want.Content)
		assert.True(got.Time.Equal(c.want.Time))
	}

	failCases := [3]uint64{0, 7, 12}

	for _, c := range failCases {
		_, err := GetOne(c)

		if err == nil {
			t.Errorf("GetOne(%v) should return an error", c)
		}
	}
}

func TestGetByID(t *testing.T) {
	assert := assert.New(t)

	want, err := GetOne(1)
	assert.Nil(err)

	got, err := GetByID(want.ID)
	assert.Nil(err)

	assert.Equal(want, got)
}

func TestEdit(t *testing.T) {
	assert := assert.New(t)

	p, _ := GetOne(3)

	err := Edit(3, "edited content")
	assert.Nil(err)

	pEdit, _ := GetOne(3)

	assert.Equal(p.ID, pEdit.ID)
	assert.Equal(p.Active, pEdit.Active)
	assert.Equal(p.Author, pEdit.Author)
	assert.Equal(pEdit.Content, "edited content")
	assert.True(p.Time.Equal(pEdit.Time))
}

func TestSubmit(t *testing.T) {
	assert := assert.New(t)

	p, _ := New("user", "content")

	err := Submit(p)

	assert.Nil(err)

	n, err := Count()
	assert.Nil(err)
	assert.Equal(n, 7)

	pt, err := GetOne(7)
	assert.Nil(err)

	assert.Equal(p.Active, pt.Active)
	assert.Equal(p.Author, pt.Author)
	assert.Equal(p.Content, pt.Content)
	assert.Equal(p.Time.Hour(), pt.Time.Hour())     // If the hour and second are equal we can be
	assert.Equal(p.Time.Second(), pt.Time.Second()) // reasonably confident the times are equal

	err = Submit(nil)

	assert.NotNil(err)
}
<commit_msg>Modify posts.Submit test for new signature<commit_after>package posts

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"testing"
	"time"

	"github.com/boatilus/peppercorn/db"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	rethink "gopkg.in/dancannon/gorethink.v2"
)

const tableName = "posts_test"

type doc struct {
	Active  bool
	Author  string
	Content string
	Time    int64
}

var docs []doc // Stores test data read in from JSON

func makePostFromDoc(d doc) Post {
	return Post{
		Active:  d.Active,
		Author:  d.Author,
		Content: d.Content,
		Time:    time.Unix(d.Time, 0),
	}
}

func setupDB() {
	if !db.Session.IsConnected() {
		panic("No DB connected")
	}

	rethink.DBCreate("peppercorn").RunWrite(db.Session)

	peppercorn := rethink.DB(db.Name)

	c, err := peppercorn.TableList().Contains(tableName).Run(db.Session)
	if err != nil {
		panic(err)
	}

	var hasTable bool

	err = c.One(&hasTable)
	if err != nil {
		panic(err)
	}

	table := peppercorn.Table(tableName)

	if !hasTable {
		_, err := peppercorn.TableCreate(tableName).RunWrite(db.Session)
		if err != nil {
			panic(err)
		}

		table.IndexCreate("active").Run(db.Session)
		table.IndexCreate("author").Run(db.Session)
		table.IndexCreate("time").Run(db.Session)
		table.IndexWait().Run(db.Session)
	} else {
		// Due to a lack of mocking in gorethink, we'll tear down the test data and repopulate on each
		// run of the tests.
		table.Delete().RunWrite(db.Session)
	}

	bytes, err := ioutil.ReadFile("posts.test_data.json")

	if err := json.Unmarshal(bytes, &docs); err != nil {
		panic(err)
	}

	if len(docs) != 7 {
		panic(err)
	}

	posts := make([]Post, len(docs))

	for i := range docs {
		posts[i].Active = docs[i].Active
		posts[i].Author = docs[i].Author
		posts[i].Content = docs[i].Content
		posts[i].Time = time.Unix(docs[i].Time, 0)
	}

	if _, err := table.Insert(posts).RunWrite(db.Session); err != nil {
		panic(err)
	}

	cursor, err := table.Count().Run(db.Session)

	if err != nil {
		panic(err)
	}

	var n int

	cursor.One(&n)
	cursor.Close()

	if n != 7 {
		panic(err)
	}
}

func init() {
	viper.Set("db.posts_table", "posts_test")

	var err error

	if db.Session, err = rethink.Connect(rethink.ConnectOpts{Address: "localhost:28015"}); err != nil {
		panic(err)
	}

	log.SetOutput(ioutil.Discard)

	setupDB()
}

///////////
// Tests //
///////////

func TestNew(t *testing.T) {
	assert := assert.New(t)

	type data struct {
		author  string
		content string
	}

	passCases := []data{
		{"user", "content"},
		{"_", "_"},
	}

	for _, c := range passCases {
		got, err := New(c.author, c.content)

		assert.Nil(err)
		assert.Empty(got.ID)
		assert.True(got.Active)
		assert.Equal(c.author, got.Author)
		assert.Equal(c.content, got.Content)
	}

	failCases := []data{
		{"", "content"},
		{"user", ""},
		{"", ""},
	}

	for _, c := range failCases {
		_, err := New(c.author, c.content)

		assert.NotNil(err)
	}
}

func TestCount(t *testing.T) {
	n, err := Count()

	assert.Nil(t, err)
	assert.Equal(t, n, 6)
}

func TestCountAll(t *testing.T) {
	n, err := CountAll()

	assert.Nil(t, err)
	assert.Equal(t, n, 7)
}

func TestGetRange(t *testing.T) {
	assert := assert.New(t)

	cases := []struct {
		first uint64
		limit uint64
		want  []Post
	}{
		{1, 2, []Post{makePostFromDoc(docs[0]), makePostFromDoc(docs[1])}},
		{3, 2, []Post{makePostFromDoc(docs[2]), makePostFromDoc(docs[3])}},
		{2, 3, []Post{makePostFromDoc(docs[1]), makePostFromDoc(docs[2]), makePostFromDoc(docs[3])}},
		// The 'first' argument is locked to 1 if < 1, so we should check that we get posts 1 and 2...
		{0, 2, []Post{makePostFromDoc(docs[0]), makePostFromDoc(docs[1])}},
	}

	for _, c := range cases {
		got, err := GetRange(c.first, c.limit)

		assert.Nil(err)

		for i := range got {
			g := got[i]
			w := c.want[i]

			assert.Equal(g.Active, w.Active)
			assert.Equal(g.Author, w.Author)
			assert.Equal(g.Content, w.Content)
			assert.True(g.Time.Equal(w.Time))
		}
	}
}

func TestGetOne(t *testing.T) {
	assert := assert.New(t)

	cases := []struct {
		in   uint64
		want Post
	}{
		{1, makePostFromDoc(docs[0])},
		{3, makePostFromDoc(docs[2])},
	}

	for _, c := range cases {
		got, err := GetOne(c.in)

		assert.Nil(err)

		assert.Equal(got.Active, c.want.Active)
		assert.Equal(got.Author, c.want.Author)
		assert.Equal(got.Content, c.want.Content)
		assert.True(got.Time.Equal(c.want.Time))
	}

	failCases := [3]uint64{0, 7, 12}

	for _, c := range failCases {
		_, err := GetOne(c)

		if err == nil {
			t.Errorf("GetOne(%v) should return an error", c)
		}
	}
}

func TestGetByID(t *testing.T) {
	assert := assert.New(t)

	want, err := GetOne(1)
	assert.Nil(err)

	got, err := GetByID(want.ID)
	assert.Nil(err)

	assert.Equal(want, got)
}

func TestEdit(t *testing.T) {
	assert := assert.New(t)

	p, _ := GetOne(3)

	err := Edit(3, "edited content")
	assert.Nil(err)

	pEdit, _ := GetOne(3)

	assert.Equal(p.ID, pEdit.ID)
	assert.Equal(p.Active, pEdit.Active)
	assert.Equal(p.Author, pEdit.Author)
	assert.Equal(pEdit.Content, "edited content")
	assert.True(p.Time.Equal(pEdit.Time))
}

func TestSubmit(t *testing.T) {
	assert := assert.New(t)

	p, _ := New("user", "content")

	id, err := Submit(p)
	assert.Nil(err)
	assert.NotEmpty(id)

	n, err := Count()
	assert.Nil(err)
	assert.Equal(n, 7)

	pt, err := GetOne(7)
	assert.Nil(err)

	assert.Equal(p.Active, pt.Active)
	assert.Equal(p.Author, pt.Author)
	assert.Equal(p.Content, pt.Content)
	assert.Equal(p.Time.Hour(), pt.Time.Hour())     // If the hour and second are equal we can be
	assert.Equal(p.Time.Second(), pt.Time.Second()) // reasonably confident the times are equal

	id, err = Submit(nil)
	assert.NotNil(err)
	assert.Empty(id)
}
<|endoftext|>


<commit_before><commit_msg>Update filter.go<commit_after><|endoftext|>