<commit_before>package queuesender

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
)

type connManagerType struct {
	refresh       func() (net.Conn, error)
	lock          sync.Mutex
	active        *connType
	connRefreshes uint64
}

// newConnManager returns a connection manager that creates connections to
// urlStr.
func newConnManager(urlStr string) (*connManagerType, error) {
	scheme, endpoint, err := extractEndpoint(urlStr)
	if err != nil {
		return nil, err
	}
	var refresh func() (net.Conn, error)
	if scheme == "https" {
		refresh = func() (net.Conn, error) {
			return tls.Dial("tcp", endpoint, nil)
		}
	}
	if refresh == nil {
		return nil, errors.New("Unsupported scheme")
	}
	return &connManagerType{refresh: refresh}, nil
}

// Get returns the current connection. If the current connection is bad,
// Get replaces it with a new connection.
func (c *connManagerType) Get() (*connType, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if c.active != nil && !c.active.isBad() {
		return c.active.get(), nil
	}
	conn, err := c.refresh()
	if err != nil {
		return nil, err
	}
	c.active = newConn(conn)
	c.connRefreshes++
	return c.active.get(), nil
}

// Refreshes returns the number of connection refreshes. That is, the
// number of times Get() created a new connection to replace a bad connection.
func (c *connManagerType) Refreshes() uint64 {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.connRefreshes
}

// connType represents a single connection.
type connType struct {
	conn      net.Conn
	br        *bufio.Reader
	lock      sync.Mutex
	userCount int
	bad       bool
}

func newConn(conn net.Conn) *connType {
	return &connType{conn: conn, br: bufio.NewReader(conn)}
}

// Put marks this connection as unused by the caller. Each time the caller
// calls Get() on the connection manager to get a connection, it must call
// Put() on the returned connection when done with it. The caller must not
// use a connection after calling Put on it.
func (c *connType) Put() {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.userCount--
	if c.userCount < 0 {
		panic("userCount cannot be negative")
	}
	if c.userCount < 1 && c.bad {
		c.conn.Close()
	}
}

// MarkBad marks this connection as bad. Caller must call MarkBad before
// calling Put.
func (c *connType) MarkBad() {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.bad = true
}

func (c *connType) isBad() bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.bad
}

func (c *connType) get() *connType {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.userCount++
	return c
}

// Send sends given request on this connection.
func (c *connType) Send(req requestType) error {
	buffer, err := encodeJSON(req.Json)
	if err != nil {
		return err
	}
	httpreq, err := http.NewRequest("POST", req.Endpoint, buffer)
	if err != nil {
		return err
	}
	httpreq.Header.Set("Content-Type", "application/json")
	return httpreq.Write(c.conn)
}

// Read reads response from the connection returning true if the connection
// is good or false if the connection is bad. If the connection is good,
// err will be non-nil if response read is something other than a 2XX
// response.
func (c *connType) Read() (bool, error) {
	resp, err := http.ReadResponse(c.br, nil)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		var buffer bytes.Buffer
		io.Copy(&buffer, resp.Body)
		// 408 is a special error meaning that the connection timed out
		// because we didn't send any request. 408 errors don't go with
		// any particular request. When we get one of these, assume the
		// connection is bad.
		return resp.StatusCode != 408, errors.New(resp.Status + ": " + buffer.String())
	}
	return true, nil
}

func encodeJSON(payload interface{}) (*bytes.Buffer, error) {
	result := &bytes.Buffer{}
	encoder := json.NewEncoder(result)
	if err := encoder.Encode(payload); err != nil {
		return nil, err
	}
	return result, nil
}

func extractEndpoint(urlStr string) (scheme, endpoint string, err error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return "", "", err
	}
	host := u.Host
	if u.Scheme == "https" && !strings.Contains(u.Host, ":") {
		host = host + ":443"
	}
	if u.Scheme == "http" && !strings.Contains(u.Host, ":") {
		host = host + ":80"
	}
	return u.Scheme, host, nil
}
<commit_msg>Support http as well as https for asynchronous communication.<commit_after>package queuesender

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
)

type connManagerType struct {
	refresh       func() (net.Conn, error)
	lock          sync.Mutex
	active        *connType
	connRefreshes uint64
}

// newConnManager returns a connection manager that creates connections to
// urlStr.
func newConnManager(urlStr string) (*connManagerType, error) {
	scheme, endpoint, err := extractEndpoint(urlStr)
	if err != nil {
		return nil, err
	}
	var refresh func() (net.Conn, error)
	if scheme == "https" {
		refresh = func() (net.Conn, error) {
			return tls.Dial("tcp", endpoint, nil)
		}
	} else if scheme == "http" {
		refresh = func() (net.Conn, error) {
			return net.Dial("tcp", endpoint)
		}
	}
	if refresh == nil {
		return nil, errors.New("Unsupported scheme")
	}
	return &connManagerType{refresh: refresh}, nil
}

// Get returns the current connection. If the current connection is bad,
// Get replaces it with a new connection.
func (c *connManagerType) Get() (*connType, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if c.active != nil && !c.active.isBad() {
		return c.active.get(), nil
	}
	conn, err := c.refresh()
	if err != nil {
		return nil, err
	}
	c.active = newConn(conn)
	c.connRefreshes++
	return c.active.get(), nil
}

// Refreshes returns the number of connection refreshes. That is, the
// number of times Get() created a new connection to replace a bad connection.
func (c *connManagerType) Refreshes() uint64 {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.connRefreshes
}

// connType represents a single connection.
type connType struct {
	conn      net.Conn
	br        *bufio.Reader
	lock      sync.Mutex
	userCount int
	bad       bool
}

func newConn(conn net.Conn) *connType {
	return &connType{conn: conn, br: bufio.NewReader(conn)}
}

// Put marks this connection as unused by the caller. Each time the caller
// calls Get() on the connection manager to get a connection, it must call
// Put() on the returned connection when done with it. The caller must not
// use a connection after calling Put on it.
func (c *connType) Put() {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.userCount--
	if c.userCount < 0 {
		panic("userCount cannot be negative")
	}
	if c.userCount < 1 && c.bad {
		c.conn.Close()
	}
}

// MarkBad marks this connection as bad. Caller must call MarkBad before
// calling Put.
func (c *connType) MarkBad() {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.bad = true
}

func (c *connType) isBad() bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.bad
}

func (c *connType) get() *connType {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.userCount++
	return c
}

// Send sends given request on this connection.
func (c *connType) Send(req requestType) error {
	buffer, err := encodeJSON(req.Json)
	if err != nil {
		return err
	}
	httpreq, err := http.NewRequest("POST", req.Endpoint, buffer)
	if err != nil {
		return err
	}
	httpreq.Header.Set("Content-Type", "application/json")
	return httpreq.Write(c.conn)
}

// Read reads response from the connection returning true if the connection
// is good or false if the connection is bad. If the connection is good,
// err will be non-nil if response read is something other than a 2XX
// response.
func (c *connType) Read() (bool, error) {
	resp, err := http.ReadResponse(c.br, nil)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		var buffer bytes.Buffer
		io.Copy(&buffer, resp.Body)
		// 408 is a special error meaning that the connection timed out
		// because we didn't send any request. 408 errors don't go with
		// any particular request. When we get one of these, assume the
		// connection is bad.
		return resp.StatusCode != 408, errors.New(resp.Status + ": " + buffer.String())
	}
	return true, nil
}

func encodeJSON(payload interface{}) (*bytes.Buffer, error) {
	result := &bytes.Buffer{}
	encoder := json.NewEncoder(result)
	if err := encoder.Encode(payload); err != nil {
		return nil, err
	}
	return result, nil
}

func extractEndpoint(urlStr string) (scheme, endpoint string, err error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return "", "", err
	}
	host := u.Host
	if u.Scheme == "https" && !strings.Contains(u.Host, ":") {
		host = host + ":443"
	}
	if u.Scheme == "http" && !strings.Contains(u.Host, ":") {
		host = host + ":80"
	}
	return u.Scheme, host, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"net/http"

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

const (
	namespace = "namenode"
)

var (
	listenAddress  = flag.String("web.listen-address", ":9070", "Address on which to expose metrics and web interface.")
	metricsPath    = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
	namenodeJmxUrl = flag.String("namenode.jmx.url", "http://localhost:50070/jmx", "Hadoop JMX URL.")
)

type Exporter struct {
	url                      string
	MissingBlocks            prometheus.Gauge
	CapacityTotal            prometheus.Gauge
	CapacityUsed             prometheus.Gauge
	CapacityRemaining        prometheus.Gauge
	CapacityUsedNonDFS       prometheus.Gauge
	BlocksTotal              prometheus.Gauge
	FilesTotal               prometheus.Gauge
	CorruptBlocks            prometheus.Gauge
	ExcessBlocks             prometheus.Gauge
	StaleDataNodes           prometheus.Gauge
	cmsGcCount               prometheus.Counter
	cmsGcTime                prometheus.Counter
	heapMemoryUsageCommitted prometheus.Gauge
	heapMemoryUsageInit      prometheus.Gauge
	heapMemoryUsageMax       prometheus.Gauge
	heapMemoryUsageUsed      prometheus.Gauge
}

func NewExporter(url string) *Exporter {
	return &Exporter{
		url: url,
		MissingBlocks: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "MissingBlocks",
			Help:      "MissingBlocks",
		}),
		CapacityTotal: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityTotal",
			Help:      "CapacityTotal",
		}),
		CapacityUsed: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityUsed",
			Help:      "CapacityUsed",
		}),
		CapacityRemaining: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityRemaining",
			Help:      "CapacityRemaining",
		}),
		CapacityUsedNonDFS: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityUsedNonDFS",
			Help:      "CapacityUsedNonDFS",
		}),
		BlocksTotal: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "BlocksTotal",
			Help:      "BlocksTotal",
		}),
		FilesTotal: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "FilesTotal",
			Help:      "FilesTotal",
		}),
		CorruptBlocks: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CorruptBlocks",
			Help:      "CorruptBlocks",
		}),
		ExcessBlocks: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "ExcessBlocks",
			Help:      "ExcessBlocks",
		}),
		StaleDataNodes: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "StaleDataNodes",
			Help:      "StaleDataNodes",
		}),
		cmsGcCount: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "ConcurrentMarkSweep_CollectionCount",
			Help:      "ConcurrentMarkSweep GC Count",
		}),
		cmsGcTime: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "ConcurrentMarkSweep_CollectionTime",
			Help:      "ConcurrentMarkSweep GC Time",
		}),
		heapMemoryUsageCommitted: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageCommitted",
			Help:      "heapMemoryUsageCommitted",
		}),
		heapMemoryUsageInit: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageInit",
			Help:      "heapMemoryUsageInit",
		}),
		heapMemoryUsageMax: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageMax",
			Help:      "heapMemoryUsageMax",
		}),
		heapMemoryUsageUsed: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageUsed",
			Help:      "heapMemoryUsageUsed",
		}),
	}
}

// Describe implements the prometheus.Collector interface.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	e.MissingBlocks.Describe(ch)
	e.CapacityTotal.Describe(ch)
	e.CapacityUsed.Describe(ch)
	e.CapacityRemaining.Describe(ch)
	e.CapacityUsedNonDFS.Describe(ch)
	e.BlocksTotal.Describe(ch)
	e.FilesTotal.Describe(ch)
	e.CorruptBlocks.Describe(ch)
	e.ExcessBlocks.Describe(ch)
	e.StaleDataNodes.Describe(ch)
	e.cmsGcCount.Describe(ch)
	e.heapMemoryUsageCommitted.Describe(ch)
	e.heapMemoryUsageInit.Describe(ch)
	e.heapMemoryUsageMax.Describe(ch)
	e.heapMemoryUsageUsed.Describe(ch)
}

// Collect implements the prometheus.Collector interface.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	resp, err := http.Get(e.url)
	if err != nil {
		log.Error(err)
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Error(err)
	}
	var f interface{}
	err = json.Unmarshal(data, &f)
	if err != nil {
		log.Error(err)
	}
	// {"beans":[{"name":"Hadoop:service=NameNode,name=FSNamesystem", ...}, {"name":"java.lang:type=MemoryPool,name=Code Cache", ...}, ...]}
	m := f.(map[string]interface{})
	// [{"name":"Hadoop:service=NameNode,name=FSNamesystem", ...}, {"name":"java.lang:type=MemoryPool,name=Code Cache", ...}, ...]
	var nameList = m["beans"].([]interface{})
	for _, nameData := range nameList {
		nameDataMap := nameData.(map[string]interface{})
		/*
			{
				"name" : "Hadoop:service=NameNode,name=FSNamesystem",
				"modelerType" : "FSNamesystem",
				"tag.Context" : "dfs",
				"tag.HAState" : "active",
				"tag.TotalSyncTimes" : "23 6 ",
				"tag.Hostname" : "CNHORTO7502.line.ism",
				"MissingBlocks" : 0,
				"MissingReplOneBlocks" : 0,
				"ExpiredHeartbeats" : 0,
				"TransactionsSinceLastCheckpoint" : 2007,
				"TransactionsSinceLastLogRoll" : 7,
				"LastWrittenTransactionId" : 172706,
				"LastCheckpointTime" : 1456089173101,
				"CapacityTotal" : 307099828224,
				"CapacityTotalGB" : 286.0,
				"CapacityUsed" : 1471291392,
				"CapacityUsedGB" : 1.0,
				"CapacityRemaining" : 279994568704,
				"CapacityRemainingGB" : 261.0,
				"CapacityUsedNonDFS" : 25633968128,
				"TotalLoad" : 6,
				"SnapshottableDirectories" : 0,
				"Snapshots" : 0,
				"LockQueueLength" : 0,
				"BlocksTotal" : 67,
				"NumFilesUnderConstruction" : 0,
				"NumActiveClients" : 0,
				"FilesTotal" : 184,
				"PendingReplicationBlocks" : 0,
				"UnderReplicatedBlocks" : 0,
				"CorruptBlocks" : 0,
				"ScheduledReplicationBlocks" : 0,
				"PendingDeletionBlocks" : 0,
				"ExcessBlocks" : 0,
				"PostponedMisreplicatedBlocks" : 0,
				"PendingDataNodeMessageCount" : 0,
				"MillisSinceLastLoadedEdits" : 0,
				"BlockCapacity" : 2097152,
				"StaleDataNodes" : 0,
				"TotalFiles" : 184,
				"TotalSyncCount" : 7
			}
		*/
		if nameDataMap["name"] == "Hadoop:service=NameNode,name=FSNamesystem" {
			e.MissingBlocks.Set(nameDataMap["MissingBlocks"].(float64))
			e.CapacityTotal.Set(nameDataMap["CapacityTotal"].(float64))
			e.CapacityUsed.Set(nameDataMap["CapacityUsed"].(float64))
			e.CapacityRemaining.Set(nameDataMap["CapacityRemaining"].(float64))
			e.CapacityUsedNonDFS.Set(nameDataMap["CapacityUsedNonDFS"].(float64))
			e.BlocksTotal.Set(nameDataMap["BlocksTotal"].(float64))
			e.FilesTotal.Set(nameDataMap["FilesTotal"].(float64))
			e.CorruptBlocks.Set(nameDataMap["CorruptBlocks"].(float64))
			e.ExcessBlocks.Set(nameDataMap["ExcessBlocks"].(float64))
			e.StaleDataNodes.Set(nameDataMap["StaleDataNodes"].(float64))
		}
		if nameDataMap["name"] == "java.lang:type=GarbageCollector,name=ConcurrentMarkSweep" {
			e.cmsGcCount.Set(nameDataMap["CollectionCount"].(float64))
			e.cmsGcTime.Set(nameDataMap["CollectionTime"].(float64))
		}
		/*
				"name" : "java.lang:type=Memory",
		    "modelerType" : "sun.management.MemoryImpl",
		    "HeapMemoryUsage" : {
		      "committed" : 1060372480,
		      "init" : 1073741824,
		      "max" : 1060372480,
		      "used" : 124571464
		    },
		*/
		if nameDataMap["name"] == "java.lang:type=Memory" {
			heapMemoryUsage := nameDataMap["HeapMemoryUsage"].(map[string]interface{})
			e.heapMemoryUsageCommitted.Set(heapMemoryUsage["committed"].(float64))
			e.heapMemoryUsageInit.Set(heapMemoryUsage["init"].(float64))
			e.heapMemoryUsageMax.Set(heapMemoryUsage["max"].(float64))
			e.heapMemoryUsageUsed.Set(heapMemoryUsage["used"].(float64))
		}

	}
	e.MissingBlocks.Collect(ch)
	e.CapacityTotal.Collect(ch)
	e.CapacityUsed.Collect(ch)
	e.CapacityRemaining.Collect(ch)
	e.CapacityUsedNonDFS.Collect(ch)
	e.BlocksTotal.Collect(ch)
	e.FilesTotal.Collect(ch)
	e.CorruptBlocks.Collect(ch)
	e.ExcessBlocks.Collect(ch)
	e.StaleDataNodes.Collect(ch)
	e.cmsGcCount.Collect(ch)
	e.cmsGcTime.Collect(ch)
	e.heapMemoryUsageCommitted.Collect(ch)
	e.heapMemoryUsageInit.Collect(ch)
	e.heapMemoryUsageMax.Collect(ch)
	e.heapMemoryUsageUsed.Collect(ch)
}

func main() {
	flag.Parse()

	exporter := NewExporter(*namenodeJmxUrl)
	prometheus.MustRegister(exporter)

	log.Printf("Starting Server: %s", *listenAddress)
	http.Handle(*metricsPath, prometheus.Handler())
	err := http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Format source<commit_after>package main

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"net/http"

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

const (
	namespace = "namenode"
)

var (
	listenAddress  = flag.String("web.listen-address", ":9070", "Address on which to expose metrics and web interface.")
	metricsPath    = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
	namenodeJmxUrl = flag.String("namenode.jmx.url", "http://localhost:50070/jmx", "Hadoop JMX URL.")
)

type Exporter struct {
	url                      string
	MissingBlocks            prometheus.Gauge
	CapacityTotal            prometheus.Gauge
	CapacityUsed             prometheus.Gauge
	CapacityRemaining        prometheus.Gauge
	CapacityUsedNonDFS       prometheus.Gauge
	BlocksTotal              prometheus.Gauge
	FilesTotal               prometheus.Gauge
	CorruptBlocks            prometheus.Gauge
	ExcessBlocks             prometheus.Gauge
	StaleDataNodes           prometheus.Gauge
	cmsGcCount               prometheus.Counter
	cmsGcTime                prometheus.Counter
	heapMemoryUsageCommitted prometheus.Gauge
	heapMemoryUsageInit      prometheus.Gauge
	heapMemoryUsageMax       prometheus.Gauge
	heapMemoryUsageUsed      prometheus.Gauge
}

func NewExporter(url string) *Exporter {
	return &Exporter{
		url: url,
		MissingBlocks: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "MissingBlocks",
			Help:      "MissingBlocks",
		}),
		CapacityTotal: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityTotal",
			Help:      "CapacityTotal",
		}),
		CapacityUsed: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityUsed",
			Help:      "CapacityUsed",
		}),
		CapacityRemaining: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityRemaining",
			Help:      "CapacityRemaining",
		}),
		CapacityUsedNonDFS: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CapacityUsedNonDFS",
			Help:      "CapacityUsedNonDFS",
		}),
		BlocksTotal: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "BlocksTotal",
			Help:      "BlocksTotal",
		}),
		FilesTotal: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "FilesTotal",
			Help:      "FilesTotal",
		}),
		CorruptBlocks: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "CorruptBlocks",
			Help:      "CorruptBlocks",
		}),
		ExcessBlocks: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "ExcessBlocks",
			Help:      "ExcessBlocks",
		}),
		StaleDataNodes: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "StaleDataNodes",
			Help:      "StaleDataNodes",
		}),
		cmsGcCount: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "ConcurrentMarkSweep_CollectionCount",
			Help:      "ConcurrentMarkSweep GC Count",
		}),
		cmsGcTime: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "ConcurrentMarkSweep_CollectionTime",
			Help:      "ConcurrentMarkSweep GC Time",
		}),
		heapMemoryUsageCommitted: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageCommitted",
			Help:      "heapMemoryUsageCommitted",
		}),
		heapMemoryUsageInit: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageInit",
			Help:      "heapMemoryUsageInit",
		}),
		heapMemoryUsageMax: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageMax",
			Help:      "heapMemoryUsageMax",
		}),
		heapMemoryUsageUsed: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "heapMemoryUsageUsed",
			Help:      "heapMemoryUsageUsed",
		}),
	}
}

// Describe implements the prometheus.Collector interface.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	e.MissingBlocks.Describe(ch)
	e.CapacityTotal.Describe(ch)
	e.CapacityUsed.Describe(ch)
	e.CapacityRemaining.Describe(ch)
	e.CapacityUsedNonDFS.Describe(ch)
	e.BlocksTotal.Describe(ch)
	e.FilesTotal.Describe(ch)
	e.CorruptBlocks.Describe(ch)
	e.ExcessBlocks.Describe(ch)
	e.StaleDataNodes.Describe(ch)
	e.cmsGcCount.Describe(ch)
	e.heapMemoryUsageCommitted.Describe(ch)
	e.heapMemoryUsageInit.Describe(ch)
	e.heapMemoryUsageMax.Describe(ch)
	e.heapMemoryUsageUsed.Describe(ch)
}

// Collect implements the prometheus.Collector interface.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	resp, err := http.Get(e.url)
	if err != nil {
		log.Error(err)
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Error(err)
	}
	var f interface{}
	err = json.Unmarshal(data, &f)
	if err != nil {
		log.Error(err)
	}
	// {"beans":[{"name":"Hadoop:service=NameNode,name=FSNamesystem", ...}, {"name":"java.lang:type=MemoryPool,name=Code Cache", ...}, ...]}
	m := f.(map[string]interface{})
	// [{"name":"Hadoop:service=NameNode,name=FSNamesystem", ...}, {"name":"java.lang:type=MemoryPool,name=Code Cache", ...}, ...]
	var nameList = m["beans"].([]interface{})
	for _, nameData := range nameList {
		nameDataMap := nameData.(map[string]interface{})
		/*
			{
				"name" : "Hadoop:service=NameNode,name=FSNamesystem",
				"modelerType" : "FSNamesystem",
				"tag.Context" : "dfs",
				"tag.HAState" : "active",
				"tag.TotalSyncTimes" : "23 6 ",
				"tag.Hostname" : "CNHORTO7502.line.ism",
				"MissingBlocks" : 0,
				"MissingReplOneBlocks" : 0,
				"ExpiredHeartbeats" : 0,
				"TransactionsSinceLastCheckpoint" : 2007,
				"TransactionsSinceLastLogRoll" : 7,
				"LastWrittenTransactionId" : 172706,
				"LastCheckpointTime" : 1456089173101,
				"CapacityTotal" : 307099828224,
				"CapacityTotalGB" : 286.0,
				"CapacityUsed" : 1471291392,
				"CapacityUsedGB" : 1.0,
				"CapacityRemaining" : 279994568704,
				"CapacityRemainingGB" : 261.0,
				"CapacityUsedNonDFS" : 25633968128,
				"TotalLoad" : 6,
				"SnapshottableDirectories" : 0,
				"Snapshots" : 0,
				"LockQueueLength" : 0,
				"BlocksTotal" : 67,
				"NumFilesUnderConstruction" : 0,
				"NumActiveClients" : 0,
				"FilesTotal" : 184,
				"PendingReplicationBlocks" : 0,
				"UnderReplicatedBlocks" : 0,
				"CorruptBlocks" : 0,
				"ScheduledReplicationBlocks" : 0,
				"PendingDeletionBlocks" : 0,
				"ExcessBlocks" : 0,
				"PostponedMisreplicatedBlocks" : 0,
				"PendingDataNodeMessageCount" : 0,
				"MillisSinceLastLoadedEdits" : 0,
				"BlockCapacity" : 2097152,
				"StaleDataNodes" : 0,
				"TotalFiles" : 184,
				"TotalSyncCount" : 7
			}
		*/
		if nameDataMap["name"] == "Hadoop:service=NameNode,name=FSNamesystem" {
			e.MissingBlocks.Set(nameDataMap["MissingBlocks"].(float64))
			e.CapacityTotal.Set(nameDataMap["CapacityTotal"].(float64))
			e.CapacityUsed.Set(nameDataMap["CapacityUsed"].(float64))
			e.CapacityRemaining.Set(nameDataMap["CapacityRemaining"].(float64))
			e.CapacityUsedNonDFS.Set(nameDataMap["CapacityUsedNonDFS"].(float64))
			e.BlocksTotal.Set(nameDataMap["BlocksTotal"].(float64))
			e.FilesTotal.Set(nameDataMap["FilesTotal"].(float64))
			e.CorruptBlocks.Set(nameDataMap["CorruptBlocks"].(float64))
			e.ExcessBlocks.Set(nameDataMap["ExcessBlocks"].(float64))
			e.StaleDataNodes.Set(nameDataMap["StaleDataNodes"].(float64))
		}
		if nameDataMap["name"] == "java.lang:type=GarbageCollector,name=ConcurrentMarkSweep" {
			e.cmsGcCount.Set(nameDataMap["CollectionCount"].(float64))
			e.cmsGcTime.Set(nameDataMap["CollectionTime"].(float64))
		}
		/*
			"name" : "java.lang:type=Memory",
			"modelerType" : "sun.management.MemoryImpl",
			"HeapMemoryUsage" : {
				"committed" : 1060372480,
				"init" : 1073741824,
				"max" : 1060372480,
				"used" : 124571464
			},
		*/
		if nameDataMap["name"] == "java.lang:type=Memory" {
			heapMemoryUsage := nameDataMap["HeapMemoryUsage"].(map[string]interface{})
			e.heapMemoryUsageCommitted.Set(heapMemoryUsage["committed"].(float64))
			e.heapMemoryUsageInit.Set(heapMemoryUsage["init"].(float64))
			e.heapMemoryUsageMax.Set(heapMemoryUsage["max"].(float64))
			e.heapMemoryUsageUsed.Set(heapMemoryUsage["used"].(float64))
		}

	}
	e.MissingBlocks.Collect(ch)
	e.CapacityTotal.Collect(ch)
	e.CapacityUsed.Collect(ch)
	e.CapacityRemaining.Collect(ch)
	e.CapacityUsedNonDFS.Collect(ch)
	e.BlocksTotal.Collect(ch)
	e.FilesTotal.Collect(ch)
	e.CorruptBlocks.Collect(ch)
	e.ExcessBlocks.Collect(ch)
	e.StaleDataNodes.Collect(ch)
	e.cmsGcCount.Collect(ch)
	e.cmsGcTime.Collect(ch)
	e.heapMemoryUsageCommitted.Collect(ch)
	e.heapMemoryUsageInit.Collect(ch)
	e.heapMemoryUsageMax.Collect(ch)
	e.heapMemoryUsageUsed.Collect(ch)
}

func main() {
	flag.Parse()

	exporter := NewExporter(*namenodeJmxUrl)
	prometheus.MustRegister(exporter)

	log.Printf("Starting Server: %s", *listenAddress)
	http.Handle(*metricsPath, prometheus.Handler())
	err := http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package rest

import (
	"fmt"
	"log"
	"reflect"
)

// TODO:
//  - Consider using/writing a more robust data validation library.
//    E.g. https://github.com/Assembli/beautiful-validity
//    This would allow for semantic validation and custom validation logic.
//    For now, we are only providing type validation.
//
//  - Make type coercion pluggable (i.e. conversion/validation of custom types).

// Rules is a collection of Rules and a reflect.Type which they correspond to.
type Rules interface {
	Contents() []*Rule
	ResourceType() reflect.Type
	Validate() error
	FilterRules(bool) Rules
}

type rules struct {
	contents     []*Rule
	resourceType reflect.Type
}

func (r rules) Contents() []*Rule {
	return r.contents
}

func (r rules) ResourceType() reflect.Type {
	return r.resourceType
}

// Validate verifies that the Rules are valid for the given reflect.Type, meaning
// they specify fields that exist and correct types. If a Rule is invalid, an
// error is returned. If the Rules are valid, nil is returned. This will
// recursively validate nested Rules.
func (r rules) Validate() error {
	for _, rule := range r.contents {
		resourceType := r.resourceType
		if resourceType.Kind() != reflect.Struct && resourceType.Kind() != reflect.Map {
			return fmt.Errorf(
				"Invalid resource type: must be struct or map, got %s",
				resourceType)
		}

		field, found := resourceType.FieldByName(rule.Field)
		if !found {
			return fmt.Errorf(
				"Invalid Rule for %s: field '%s' does not exist",
				resourceType, rule.Field)
		}

		if !rule.validType(field.Type) {
			return fmt.Errorf(
				"Invalid Rule for %s: field '%s' is type %s, not %s",
				resourceType, rule.Field, field.Type, typeToName[rule.Type])
		}

		// Validate nested Rules.
		if rule.Rules != nil {
			if err := rule.Rules.Validate(); err != nil {
				return err
			}
		}
	}

	return nil
}

// FilterRules filters the slice of Rules based on the specified bool. True means to
// filter out outbound Rules such that the returned slice contains only inbound Rules.
// False means to filter out inbound Rules such that the returned slice contains only
// outbound Rules.
func (r rules) FilterRules(inbound bool) Rules {
	filtered := make([]*Rule, 0, len(r.contents))
	for _, rule := range r.contents {
		if inbound && rule.OutputOnly {
			// Filter out outbound Rules.
			continue
		} else if !inbound && rule.InputOnly {
			// Filter out inbound Rules.
			continue
		}
		filtered = append(filtered, rule)
	}

	return rules{contents: filtered, resourceType: r.resourceType}
}

// NewRules returns a set of Rules for use by a ResourceHandler. The first argument
// must be a resource pointer (and can be nil) used to associate the Rules with a
// resource type. If it isn't a pointer, this will panic.
func NewRules(ptr interface{}, r ...*Rule) Rules {
	resourceType := reflect.TypeOf(ptr)
	if resourceType.Kind() != reflect.Ptr {
		panic(fmt.Sprintf("Must provide resource pointer to NewRules, got %s",
			resourceType.Kind()))
	}

	return rules{
		resourceType: resourceType.Elem(),
		contents:     r,
	}
}

// Rule provides schema validation and type coercion for request input and fine-grained
// control over response output. If a ResourceHandler provides input Rules which
// specify types, input fields will attempt to be coerced to those types. If coercion
// fails, an error will be returned in the response. If a ResourceHandler provides
// output Rules, only the fields corresponding to those Rules will be sent back. This
// prevents new fields from leaking into old API versions.
type Rule struct {
	// Name of the resource field. This is the name as it appears in the struct
	// definition.
	Field string

	// Name of the input/output field. Use Name() to retrieve the field alias while
	// falling back to the field name if it's not specified.
	FieldAlias string

	// Type to coerce field value to. If the value cannot be coerced, an error will be
	// returned in the response. Defaults to Unspecified, which is the equivalent of
	// an interface{} value.
	Type Type

	// Indicates if the field must have a value. Defaults to false.
	Required bool

	// Versions is a list of the API versions this Rule applies to. If empty, it will
	// be applied to all versions.
	Versions []string

	// Indicates if the Rule should only be applied to requests.
	InputOnly bool

	// Indicates if the Rule should only be applied to responses.
	OutputOnly bool

	// Function which produces the field value to receive.
	InputHandler func(interface{}) interface{}

	// Function which produces the field value to send.
	OutputHandler func(interface{}) interface{}

	// Nested Rules to apply to field value.
	Rules Rules
}

// Name returns the name of the input/output field alias. It defaults to the field
// name if the alias was not specified.
func (r Rule) Name() string {
	alias := r.FieldAlias
	if alias == "" {
		alias = r.Field
	}
	return alias
}

// Applies returns whether or not the Rule applies to the given version.
func (r Rule) Applies(version string) bool {
	if r.Versions == nil {
		return true
	}

	for _, v := range r.Versions {
		if v == version {
			return true
		}
	}

	return false
}

// validType returns whether or not the Rule is valid for the given reflect.Type.
func (r Rule) validType(fieldType reflect.Type) bool {
	if r.Type == Unspecified {
		return true
	}

	kind := typeToKind[r.Type]
	return fieldType.Kind() == kind
}

// applyInboundRules applies Rules which are not specified as output only to the
// provided Payload. If the Payload is nil, an empty Payload will be returned. If no
// Rules are provided, this acts as an identity function. If Rules are provided, any
// incoming fields which are not specified will be discarded. If Rules specify types,
// incoming values will attempted to be coerced. If coercion fails, an error will be
// returned.
func applyInboundRules(payload Payload, rules Rules) (Payload, error) {
	if payload == nil {
		return Payload{}, nil
	}

	// Apply only inbound Rules.
	rules = rules.FilterRules(true)

	if len(rules.Contents()) == 0 {
		return payload, nil
	}

	newPayload := Payload{}

fieldLoop:
	for field, value := range payload {
		for _, rule := range rules.Contents() {
			if rule.FieldAlias == field {
				if rule.Type != Unspecified {
					// Coerce to specified type.
					coerced, err := coerceType(value, rule.Type)
					if err != nil {
						return nil, err
					}
					value = coerced
				}
				if rule.InputHandler != nil {
					value = rule.InputHandler(value)
				}
				newPayload[field] = value
				continue fieldLoop
			}
		}

		log.Printf("Discarding field '%s'", field)
	}

	// Ensure no required fields are missing.
	if err := enforceRequiredFields(rules, newPayload); err != nil {
		log.Println(err)
		return nil, err
	}

	return newPayload, nil
}

// applyOutboundRules applies Rules which are not specified as input only to the
// provided Resource. If the Resource is nil, not a struct or
// map[string]interface{}, or no Rules are provided, this acts as an identity
// function. If Rules are provided, only the fields specified by them will be
// included in the returned Resource. This is to prevent new fields from leaking
// into old API versions.
func applyOutboundRules(resource Resource, rules Rules) Resource {
	// Apply only outbound Rules.
	rules = rules.FilterRules(false)

	if isNil(resource) || len(rules.Contents()) == 0 {
		// Return resource as-is if no Rules are provided.
		return resource
	}

	// Get the underlying value by dereferencing the pointer if there is one.
	resourceValue := reflect.Indirect(reflect.ValueOf(resource))
	resource = resourceValue.Interface()
	resourceType := reflect.TypeOf(resource)
	var payload Resource

	if resourceType.Kind() == reflect.Map {
		if resourceMap, ok := resource.(map[string]interface{}); ok {
			payload = applyOutboundRulesForMap(resourceMap, rules)
		} else {
			// Nothing we can do if the keys aren't strings.
			payload = resource
		}
	} else if resourceType.Kind() == reflect.Struct {
		payload = applyOutboundRulesForStruct(resourceValue, rules)
	} else {
		// Only apply Rules to resource structs and maps.
		payload = resource
	}

	return payload
}

// applyOutboundRulesForMap applies Rules which are not specified as input only to the
// provided map. If a Rule specifies a field which is not in the map, it will be skipped.
// If a Rule specifies nested Rules, they will be recursively applied to the corresponding
// value.
func applyOutboundRulesForMap(resource map[string]interface{}, rules Rules) Payload {
	payload := Payload{}
	for _, rule := range rules.Contents() {
		fieldValue, ok := resource[rule.Field]
		if !ok {
			log.Printf("Map resource missing field '%s'", rule.Field)
			continue
		}

		if rule.Rules != nil {
			fieldValue = applyNestedOutboundRules(fieldValue, rule)
		}

		if rule.OutputHandler != nil {
			fieldValue = rule.OutputHandler(fieldValue)
		}
		payload[rule.Name()] = fieldValue
	}

	return payload
}

// applyOutboundRulesForStruct applies Rules which are not specified as input only to the
// provided reflect.Value. The precondition for this function is that the value is an
// instance of the type specified on the Rules. If a Rule specifies nested Rules, they
// will be recursively applied to the corresponding value.
func applyOutboundRulesForStruct(resourceValue reflect.Value, rules Rules) Payload {
	payload := Payload{}
	for _, rule := range rules.Contents() {
		// Rule validation occurs at server start. No need to check for field existence.
		field := resourceValue.FieldByName(rule.Field)
		fieldValue := field.Interface()

		if rule.Rules != nil {
			fieldValue = applyNestedOutboundRules(fieldValue, rule)
		}

		if rule.OutputHandler != nil {
			fieldValue = rule.OutputHandler(fieldValue)
		}
		payload[rule.Name()] = fieldValue
	}

	return payload
}

// applyNestedOutboundRules recursively applies nested Rules which are not specified as
// input only to the provided Resource.
func applyNestedOutboundRules(resource Resource, rule *Rule) Resource {
	var fieldValue Resource

	if reflect.TypeOf(resource).Kind() == reflect.Slice {
		// Apply nested Rules to each item in the slice.
		s := reflect.ValueOf(resource)
		nestedValues := make([]interface{}, s.Len())
		for i := 0; i < s.Len(); i++ {
			nestedValues[i] = applyOutboundRules(s.Index(i).Interface(), rule.Rules)
		}
		fieldValue = nestedValues
	} else {
		fieldValue = applyOutboundRules(resource, rule.Rules)
	}

	return fieldValue
}

// enforceRequiredFields verifies that the provided Payload has values for any Rules
// with the Required flag set to true. If any required fields are missing, an error
// will be returned. Otherwise nil is returned.
func enforceRequiredFields(rules Rules, payload Payload) error {
ruleLoop:
	for _, rule := range rules.Contents() {
		if !rule.Required {
			continue
		}

		for field := range payload {
			if rule.Name() == field {
				continue ruleLoop
			}
		}

		return fmt.Errorf("Missing required field '%s'", rule.Name())
	}

	return nil
}

// isNil returns true if the given Resource is a nil value or pointer, false if
// not.
func isNil(resource Resource) bool {
	value := reflect.ValueOf(resource)
	t := value.Kind()
	switch t {
	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map,
		reflect.Ptr, reflect.Slice:
		return value.IsNil()
	default:
		return resource == nil
	}
}
<commit_msg>Apply inbound nested Rules<commit_after>package rest

import (
	"fmt"
	"log"
	"reflect"
)

// TODO:
//  - Consider using/writing a more robust data validation library.
//    E.g. https://github.com/Assembli/beautiful-validity
//    This would allow for semantic validation and custom validation logic.
//    For now, we are only providing type validation.
//
//  - Make type coercion pluggable (i.e. conversion/validation of custom types).

// Rules is a collection of Rules and a reflect.Type which they correspond to.
type Rules interface {
	// Contents returns the contained Rules.
	Contents() []*Rule

	// ResourceType returns the reflect.Type these Rules correspond to.
	ResourceType() reflect.Type

	// Validate verifies that the Rules are valid, meaning they specify fields that exist
	// and correct types. If a Rule is invalid, an error is returned. If the Rules are
	// valid, nil is returned. This will recursively validate nested Rules.
	Validate() error

	// Filter filters the slice of Rules based on the specified bool. True means to
	// filter out outbound Rules such that the returned slice contains only inbound Rules.
	// False means to filter out inbound Rules such that the returned slice contains only
	// outbound Rules.
	Filter(bool) Rules

	// Size returns the number of contained Rules.
	Size() int
}

type rules struct {
	contents     []*Rule
	resourceType reflect.Type
}

// Contents returns the contained Rules.
func (r rules) Contents() []*Rule {
	return r.contents
}

// ResourceType returns the reflect.Type these Rules correspond to.
func (r rules) ResourceType() reflect.Type {
	return r.resourceType
}

// Validate verifies that the Rules are valid, meaning they specify fields that exist
// and correct types. If a Rule is invalid, an error is returned. If the Rules are
// valid, nil is returned. This will recursively validate nested Rules.
func (r rules) Validate() error {
	for _, rule := range r.contents {
		resourceType := r.resourceType
		if resourceType.Kind() != reflect.Struct && resourceType.Kind() != reflect.Map {
			return fmt.Errorf(
				"Invalid resource type: must be struct or map, got %s",
				resourceType)
		}

		field, found := resourceType.FieldByName(rule.Field)
		if !found {
			return fmt.Errorf(
				"Invalid Rule for %s: field '%s' does not exist",
				resourceType, rule.Field)
		}

		if !rule.validType(field.Type) {
			return fmt.Errorf(
				"Invalid Rule for %s: field '%s' is type %s, not %s",
				resourceType, rule.Field, field.Type, typeToName[rule.Type])
		}

		// Validate nested Rules.
		if rule.Rules != nil {
			if err := rule.Rules.Validate(); err != nil {
				return err
			}
		}
	}

	return nil
}

// Filter filters the slice of Rules based on the specified bool. True means to
// filter out outbound Rules such that the returned slice contains only inbound Rules.
// False means to filter out inbound Rules such that the returned slice contains only
// outbound Rules.
func (r rules) Filter(inbound bool) Rules {
	filtered := make([]*Rule, 0, len(r.contents))
	for _, rule := range r.contents {
		if inbound && rule.OutputOnly {
			// Filter out outbound Rules.
			continue
		} else if !inbound && rule.InputOnly {
			// Filter out inbound Rules.
			continue
		}
		filtered = append(filtered, rule)
	}

	return rules{contents: filtered, resourceType: r.resourceType}
}

// Size returns the number of contained Rules.
func (r rules) Size() int {
	return len(r.contents)
}

// NewRules returns a set of Rules for use by a ResourceHandler. The first argument
// must be a resource pointer (and can be nil) used to associate the Rules with a
// resource type. If it isn't a pointer, this will panic.
func NewRules(ptr interface{}, r ...*Rule) Rules {
	resourceType := reflect.TypeOf(ptr)
	if resourceType.Kind() != reflect.Ptr {
		panic(fmt.Sprintf("Must provide resource pointer to NewRules, got %s",
			resourceType.Kind()))
	}

	return rules{
		resourceType: resourceType.Elem(),
		contents:     r,
	}
}

// Rule provides schema validation and type coercion for request input and fine-grained
// control over response output. If a ResourceHandler provides input Rules which
// specify types, input fields will attempt to be coerced to those types. If coercion
// fails, an error will be returned in the response. If a ResourceHandler provides
// output Rules, only the fields corresponding to those Rules will be sent back. This
// prevents new fields from leaking into old API versions.
type Rule struct {
	// Name of the resource field. This is the name as it appears in the struct
	// definition.
	Field string

	// Name of the input/output field. Use Name() to retrieve the field alias while
	// falling back to the field name if it's not specified.
	FieldAlias string

	// Type to coerce field value to. If the value cannot be coerced, an error will be
	// returned in the response. Defaults to Unspecified, which is the equivalent of
	// an interface{} value.
	Type Type

	// Indicates if the field must have a value. Defaults to false.
	Required bool

	// Versions is a list of the API versions this Rule applies to. If empty, it will
	// be applied to all versions.
	Versions []string

	// Indicates if the Rule should only be applied to requests.
	InputOnly bool

	// Indicates if the Rule should only be applied to responses.
	OutputOnly bool

	// Function which produces the field value to receive.
	InputHandler func(interface{}) interface{}

	// Function which produces the field value to send.
	OutputHandler func(interface{}) interface{}

	// Nested Rules to apply to field value.
	Rules Rules
}

// Name returns the name of the input/output field alias. It defaults to the field
// name if the alias was not specified.
func (r Rule) Name() string {
	alias := r.FieldAlias
	if alias == "" {
		alias = r.Field
	}
	return alias
}

// Applies returns whether or not the Rule applies to the given version.
func (r Rule) Applies(version string) bool {
	if r.Versions == nil {
		return true
	}

	for _, v := range r.Versions {
		if v == version {
			return true
		}
	}

	return false
}

// validType returns whether or not the Rule is valid for the given reflect.Type.
func (r Rule) validType(fieldType reflect.Type) bool {
	if r.Type == Unspecified {
		return true
	}

	kind := typeToKind[r.Type]
	return fieldType.Kind() == kind
}

// applyInboundRules applies Rules which are not specified as output only to the
// provided Payload. If the Payload is nil, an empty Payload will be returned. If no
// Rules are provided, this acts as an identity function. If Rules are provided, any
// incoming fields which are not specified will be discarded. If Rules specify types,
// incoming values will attempted to be coerced. If coercion fails, an error will be
// returned.
func applyInboundRules(payload Payload, rules Rules) (Payload, error) {
	if payload == nil {
		return Payload{}, nil
	}

	// Apply only inbound Rules.
	rules = rules.Filter(true)

	if len(rules.Contents()) == 0 {
		return payload, nil
	}

	newPayload := Payload{}

fieldLoop:
	for field, value := range payload {
		for _, rule := range rules.Contents() {
			if rule.FieldAlias == field {
				if nestedInboundRulesApply(value, rule.Rules) {
					// Nested Rules take precedence over type coercion.
					v, err := applyNestedInboundRules(value, rule.Rules)
					if err != nil {
						return nil, err
					}
					value = v
				} else if rule.Type != Unspecified {
					// Coerce to specified type.
					coerced, err := coerceType(value, rule.Type)
					if err != nil {
						return nil, err
					}
					value = coerced
				}

				if rule.InputHandler != nil {
					value = rule.InputHandler(value)
				}

				newPayload[field] = value
				continue fieldLoop
			}
		}

		log.Printf("Discarding field '%s'", field)
	}

	// Ensure no required fields are missing.
	if err := enforceRequiredFields(rules, newPayload); err != nil {
		log.Println(err)
		return nil, err
	}

	return newPayload, nil
}

// applyNestedInboundRules recursively applies nested Rules which are not specified as
// output only to the provided value.
func applyNestedInboundRules(value interface{}, rules Rules) (interface{}, error) {
	var fieldValue interface{}
	valueType := reflect.TypeOf(value).Kind()
	if valueType == reflect.Slice {
		// Apply nested Rules to each item in the slice.
		s := reflect.ValueOf(value)
		nestedValues := make([]interface{}, s.Len())
		for i := 0; i < s.Len(); i++ {
			var payload map[string]interface{}
			payload, err := applyInboundRules(
				s.Index(i).Interface().(map[string]interface{}), rules)
			if err != nil {
				return nil, err
			}
			nestedValues[i] = payload
		}
		fieldValue = nestedValues
	} else {
		var payload map[string]interface{}
		payload, err := applyInboundRules(value.(map[string]interface{}), rules)
		if err != nil {
			return nil, err
		}
		fieldValue = payload
	}

	return fieldValue, nil
}

// nestedInboundRulesApply returns true if the Rules contain inbound Rules and
// the value is a map or slice.
func nestedInboundRulesApply(value interface{}, rules Rules) bool {
	if rules == nil || rules.Size() == 0 {
		return false
	}

	valueType := reflect.TypeOf(value).Kind()
	if valueType != reflect.Map && valueType != reflect.Slice {
		// Only apply nested Rules to maps and slices.
		return false
	}

	return rules.Filter(true).Size() > 0
}

// applyOutboundRules applies Rules which are not specified as input only to the
// provided Resource. If the Resource is nil, not a struct or
// map[string]interface{}, or no Rules are provided, this acts as an identity
// function. If Rules are provided, only the fields specified by them will be
// included in the returned Resource. This is to prevent new fields from leaking
// into old API versions. If Rules specify nested Rules, they will be recursively
// applied to field values.
func applyOutboundRules(resource Resource, rules Rules) Resource {
	// Apply only outbound Rules.
	rules = rules.Filter(false)

	if isNil(resource) || len(rules.Contents()) == 0 {
		// Return resource as-is if no Rules are provided.
		return resource
	}

	// Get the underlying value by dereferencing the pointer if there is one.
	resourceValue := reflect.Indirect(reflect.ValueOf(resource))
	resource = resourceValue.Interface()
	resourceType := reflect.TypeOf(resource)
	var payload Resource

	if resourceType.Kind() == reflect.Map {
		if resourceMap, ok := resource.(map[string]interface{}); ok {
			payload = applyOutboundRulesForMap(resourceMap, rules)
		} else {
			// Nothing we can do if the keys aren't strings.
			payload = resource
		}
	} else if resourceType.Kind() == reflect.Struct {
		payload = applyOutboundRulesForStruct(resourceValue, rules)
	} else {
		// Only apply Rules to resource structs and maps.
		payload = resource
	}

	return payload
}

// applyOutboundRulesForMap applies Rules which are not specified as input only to the
// provided map. If a Rule specifies a field which is not in the map, it will be skipped.
// If a Rule specifies nested Rules, they will be recursively applied to the corresponding
// value.
func applyOutboundRulesForMap(resource map[string]interface{}, rules Rules) Payload {
	payload := Payload{}
	for _, rule := range rules.Contents() {
		fieldValue, ok := resource[rule.Field]
		if !ok {
			log.Printf("Map resource missing field '%s'", rule.Field)
			continue
		}

		if rule.Rules != nil {
			fieldValue = applyNestedOutboundRules(fieldValue, rule)
		}

		if rule.OutputHandler != nil {
			fieldValue = rule.OutputHandler(fieldValue)
		}
		payload[rule.Name()] = fieldValue
	}

	return payload
}

// applyOutboundRulesForStruct applies Rules which are not specified as input only to the
// provided reflect.Value. The precondition for this function is that the value is an
// instance of the type specified on the Rules. If a Rule specifies nested Rules, they
// will be recursively applied to the corresponding value.
func applyOutboundRulesForStruct(resourceValue reflect.Value, rules Rules) Payload {
	payload := Payload{}
	for _, rule := range rules.Contents() {
		// Rule validation occurs at server start. No need to check for field existence.
		field := resourceValue.FieldByName(rule.Field)
		fieldValue := field.Interface()

		if rule.Rules != nil {
			fieldValue = applyNestedOutboundRules(fieldValue, rule)
		}

		if rule.OutputHandler != nil {
			fieldValue = rule.OutputHandler(fieldValue)
		}
		payload[rule.Name()] = fieldValue
	}

	return payload
}

// applyNestedOutboundRules recursively applies nested Rules which are not specified as
// input only to the provided Resource.
func applyNestedOutboundRules(resource Resource, rule *Rule) Resource {
	var fieldValue Resource

	if reflect.TypeOf(resource).Kind() == reflect.Slice {
		// Apply nested Rules to each item in the slice.
		s := reflect.ValueOf(resource)
		nestedValues := make([]interface{}, s.Len())
		for i := 0; i < s.Len(); i++ {
			nestedValues[i] = applyOutboundRules(s.Index(i).Interface(), rule.Rules)
		}
		fieldValue = nestedValues
	} else {
		fieldValue = applyOutboundRules(resource, rule.Rules)
	}

	return fieldValue
}

// enforceRequiredFields verifies that the provided Payload has values for any Rules
// with the Required flag set to true. If any required fields are missing, an error
// will be returned. Otherwise nil is returned.
func enforceRequiredFields(rules Rules, payload Payload) error {
ruleLoop:
	for _, rule := range rules.Contents() {
		if !rule.Required {
			continue
		}

		for field := range payload {
			if rule.Name() == field {
				continue ruleLoop
			}
		}

		return fmt.Errorf("Missing required field '%s'", rule.Name())
	}

	return nil
}

// isNil returns true if the given Resource is a nil value or pointer, false if
// not.
func isNil(resource Resource) bool {
	value := reflect.ValueOf(resource)
	t := value.Kind()
	switch t {
	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map,
		reflect.Ptr, reflect.Slice:
		return value.IsNil()
	default:
		return resource == nil
	}
}
<|endoftext|>


<commit_before>// +build riotweb

package embed

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

	"github.com/gorilla/mux"
	"github.com/tidwall/sjson"
)

// From within the Riot Web directory:
// go run github.com/mjibson/esc -o /path/to/dendrite/internal/embed/fs_riotweb.go -private -pkg embed .

func Embed(rootMux *mux.Router, listenPort int, serverName string) {
	url := fmt.Sprintf("http://localhost:%d", listenPort)
	embeddedFS := _escFS(false)
	embeddedServ := http.FileServer(embeddedFS)

	rootMux.NotFoundHandler = embeddedServ
	rootMux.HandleFunc("/config.json", func(w http.ResponseWriter, _ *http.Request) {
		configFile, err := embeddedFS.Open("/config.sample.json")
		if err != nil {
			w.WriteHeader(500)
			io.WriteString(w, "Couldn't open the file: "+err.Error())
			return
		}
		configFileInfo, err := configFile.Stat()
		if err != nil {
			w.WriteHeader(500)
			io.WriteString(w, "Couldn't stat the file: "+err.Error())
			return
		}
		buf := make([]byte, configFileInfo.Size())
		n, err := configFile.Read(buf)
		if err != nil {
			w.WriteHeader(500)
			io.WriteString(w, "Couldn't read the file: "+err.Error())
			return
		}
		if int64(n) != configFileInfo.Size() {
			w.WriteHeader(500)
			io.WriteString(w, "The returned file size didn't match what we expected")
			return
		}
		js, _ := sjson.SetBytes(buf, "default_server_config.m\\.homeserver.base_url", url)
		js, _ = sjson.SetBytes(js, "default_server_config.m\\.homeserver.server_name", serverName)
		js, _ = sjson.SetBytes(js, "brand", fmt.Sprintf("Riot %s", serverName))
		js, _ = sjson.SetBytes(js, "disable_guests", true)
		js, _ = sjson.SetBytes(js, "disable_3pid_login", true)
		js, _ = sjson.DeleteBytes(js, "welcomeUserId")
		_, _ = w.Write(js)
	})

	fmt.Println("*-------------------------------*")
	fmt.Println("| This build includes Riot Web! |")
	fmt.Println("*-------------------------------*")
	fmt.Println("Point your browser to:", url)
	fmt.Println()
}
<commit_msg>Fix MIME types for embedded Riot in Ygg demo<commit_after>// +build riotweb

package embed

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

	"github.com/gorilla/mux"
	"github.com/tidwall/sjson"
)

// From within the Riot Web directory:
// go run github.com/mjibson/esc -o /path/to/dendrite/internal/embed/fs_riotweb.go -private -pkg embed .

var cssFile = regexp.MustCompile("\\.css$")
var jsFile = regexp.MustCompile("\\.js$")

type mimeFixingHandler struct {
	fs http.Handler
}

func (h mimeFixingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ruri := r.RequestURI
	fmt.Println(ruri)
	switch {
	case cssFile.MatchString(ruri):
		w.Header().Set("Content-Type", "text/css")
	case jsFile.MatchString(ruri):
		w.Header().Set("Content-Type", "application/javascript")
	default:
	}
	h.fs.ServeHTTP(w, r)
}

func Embed(rootMux *mux.Router, listenPort int, serverName string) {
	url := fmt.Sprintf("http://localhost:%d", listenPort)
	embeddedFS := _escFS(false)
	embeddedServ := mimeFixingHandler{http.FileServer(embeddedFS)}

	rootMux.NotFoundHandler = embeddedServ
	rootMux.HandleFunc("/config.json", func(w http.ResponseWriter, _ *http.Request) {
		configFile, err := embeddedFS.Open("/config.sample.json")
		if err != nil {
			w.WriteHeader(500)
			io.WriteString(w, "Couldn't open the file: "+err.Error())
			return
		}
		configFileInfo, err := configFile.Stat()
		if err != nil {
			w.WriteHeader(500)
			io.WriteString(w, "Couldn't stat the file: "+err.Error())
			return
		}
		buf := make([]byte, configFileInfo.Size())
		n, err := configFile.Read(buf)
		if err != nil {
			w.WriteHeader(500)
			io.WriteString(w, "Couldn't read the file: "+err.Error())
			return
		}
		if int64(n) != configFileInfo.Size() {
			w.WriteHeader(500)
			io.WriteString(w, "The returned file size didn't match what we expected")
			return
		}
		js, _ := sjson.SetBytes(buf, "default_server_config.m\\.homeserver.base_url", url)
		js, _ = sjson.SetBytes(js, "default_server_config.m\\.homeserver.server_name", serverName)
		js, _ = sjson.SetBytes(js, "brand", fmt.Sprintf("Riot %s", serverName))
		js, _ = sjson.SetBytes(js, "disable_guests", true)
		js, _ = sjson.SetBytes(js, "disable_3pid_login", true)
		js, _ = sjson.DeleteBytes(js, "welcomeUserId")
		_, _ = w.Write(js)
	})

	fmt.Println("*-------------------------------*")
	fmt.Println("| This build includes Riot Web! |")
	fmt.Println("*-------------------------------*")
	fmt.Println("Point your browser to:", url)
	fmt.Println()
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/golang/protobuf/jsonpb"
)

var marshaler = &jsonpb.Marshaler{
	EnumAsInts:   false,
	EmitDefaults: true,
	Indent:       "  ",
	OrigName:     true,
}

func main() {
}
<commit_msg>skeleton<commit_after>package main

import (
	"net/http"

	"github.com/golang/protobuf/jsonpb"
	"github.com/gorilla/pat"

	pb "github.com/tiziano88/gelm/server/proto"
)

var marshaler = &jsonpb.Marshaler{
	EnumsAsInts:  false,
	EmitDefaults: true,
	Indent:       "  ",
	OrigName:     true,
}

func main() {
	m := pat.New()

	m.Get("/", Handler)

	http.Handle("/", m)
}

func Handler(w http.ResponseWriter, r *http.Request) {
	m := &pb.Message{}
	marshaler.Marshal(w, m)
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"io"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pps"
)

// NewJob creates a pps.Job.
func NewJob(jobID string) *pps.Job {
	return &pps.Job{ID: jobID}
}

const (
	// PPSEtcdPrefixEnv is the environment variable that specifies the etcd
	// prefix that PPS uses.
	PPSEtcdPrefixEnv = "PPS_ETCD_PREFIX"
	// PPSWorkerIPEnv is the environment variable that a worker can use to
	// see its own IP.  The IP address is made available through the
	// Kubernetes downward API.
	PPSWorkerIPEnv = "PPS_WORKER_IP"
	// PPSPipelineNameEnv is the env var that sets the name of the pipeline
	// that the workers are running.
	PPSPipelineNameEnv = "PPS_PIPELINE_NAME"
	// PPSJobIDEnv is the env var that sets the ID of the job that the
	// workers are running (if the workers belong to an orphan job, rather than a
	// pipeline).
	PPSJobIDEnv = "PPS_JOB_ID"
	// PPSInputPrefix is the prefix of the path where datums are downloaded
	// to.  A datum of an input named `XXX` is downloaded to `/pfs/XXX/`.
	PPSInputPrefix = "/pfs"
	// PPSOutputPath is the path where the user code is
	// expected to write its output to.
	PPSOutputPath = "/pfs/out"
	// PPSWorkerPort is the port that workers use for their gRPC server
	PPSWorkerPort = 30652
)

// NewJobInput creates a pps.JobInput.
func NewJobInput(repoName string, commitID string, glob string) *pps.JobInput {
	return &pps.JobInput{
		Commit: NewCommit(repoName, commitID),
		Glob:   glob,
	}
}

// NewPipeline creates a pps.Pipeline.
func NewPipeline(pipelineName string) *pps.Pipeline {
	return &pps.Pipeline{Name: pipelineName}
}

// NewPipelineInput creates a new pps.PipelineInput
func NewPipelineInput(repoName string, glob string) *pps.PipelineInput {
	return &pps.PipelineInput{
		Repo: NewRepo(repoName),
		Glob: glob,
	}
}

// CreateJob creates and runs a job in PPS.
// image is the Docker image to run the job in.
// cmd is the command passed to the Docker run invocation.
// NOTE as with Docker cmd is not run inside a shell that means that things
// like wildcard globbing (*), pipes (|) and file redirects (> and >>) will not
// work. To get that behavior you should have your command be a shell of your
// choice and pass a shell script to stdin.
// stdin is a slice of lines that are sent to your command on stdin. Lines need
// not end in newline characters.
// parallelism is how many copies of your container should run in parallel. You
// may pass 0 for parallelism in which case PPS will set the parallelism based
// on availabe resources.
// inputs specifies a set of Commits that will be visible to the job during runtime.
// parentJobID specifies the a job to use as a parent, it may be left empty in
// which case there is no parent job. If not left empty your job will use the
// parent Job's output commit as the parent of its output commit.
func (c APIClient) CreateJob(
	image string,
	cmd []string,
	stdin []string,
	parallelismSpec *pps.ParallelismSpec,
	inputs []*pps.JobInput,
	internalPort int32,
	externalPort int32,
) (*pps.Job, error) {
	var service *pps.Service
	if internalPort != 0 {
		service = &pps.Service{
			InternalPort: internalPort,
		}
	}
	if externalPort != 0 {
		if internalPort == 0 {
			return nil, fmt.Errorf("external port specified without internal port")
		}
		service.ExternalPort = externalPort
	}
	job, err := c.PpsAPIClient.CreateJob(
		c.ctx(),
		&pps.CreateJobRequest{
			Transform: &pps.Transform{
				Image: image,
				Cmd:   cmd,
				Stdin: stdin,
			},
			ParallelismSpec: parallelismSpec,
			Inputs:          inputs,
			Service:         service,
		},
	)
	return job, sanitizeErr(err)
}

// InspectJob returns info about a specific job.
// blockOutput will cause the call to block until the job has been assigned an output commit.
// blockState will cause the call to block until the job reaches a terminal state (failure or success).
func (c APIClient) InspectJob(jobID string, blockState bool) (*pps.JobInfo, error) {
	jobInfo, err := c.PpsAPIClient.InspectJob(
		c.ctx(),
		&pps.InspectJobRequest{
			Job:        NewJob(jobID),
			BlockState: blockState,
		})
	return jobInfo, sanitizeErr(err)
}

// ListJob returns info about all jobs.
// If pipelineName is non empty then only jobs that were started by the named pipeline will be returned
// If inputCommit is non-nil then only jobs which took the specific commits as inputs will be returned.
// The order of the inputCommits doesn't matter.
func (c APIClient) ListJob(pipelineName string, inputCommit []*pfs.Commit) ([]*pps.JobInfo, error) {
	var pipeline *pps.Pipeline
	if pipelineName != "" {
		pipeline = NewPipeline(pipelineName)
	}
	jobInfos, err := c.PpsAPIClient.ListJob(
		c.ctx(),
		&pps.ListJobRequest{
			Pipeline:    pipeline,
			InputCommit: inputCommit,
		})
	if err != nil {
		return nil, sanitizeErr(err)
	}
	return jobInfos.JobInfo, nil
}

// DeleteJob deletes a job along with its output Repo
func (c APIClient) DeleteJob(jobID string) error {
	_, err := c.PpsAPIClient.DeleteJob(
		c.ctx(),
		&pps.DeleteJobRequest{
			Job: NewJob(jobID),
		},
	)
	return sanitizeErr(err)
}

// LogsIter iterates through log messages returned from pps.GetLogs. Logs can
// be fetched with 'Next()'. The log message received can be examined with
// 'Message()', and any errors can be examined with 'Err()'.
type LogsIter struct {
	logsClient pps.API_GetLogsClient
	msg        *pps.LogMessage
	err        error
}

func (l *LogsIter) Next() bool {
	if l.err != nil {
		l.msg = nil
		return false
	}
	l.msg, l.err = l.logsClient.Recv()
	if l.err != nil {
		return false
	}
	return true
}

func (l *LogsIter) Message() *pps.LogMessage {
	return l.msg
}

func (l *LogsIter) Err() error {
	if l.err == io.EOF {
		return nil
	}
	return l.err
}

// GetLogs gets logs from a job (logs includes stdout and stderr). 'pipelineName',
// 'jobId', and 'data', are all filters. To forego any filter, simply pass an
// empty value, though one of 'pipelineName' and 'jobId' must be set. Responses
// are written to 'messages'
func (c APIClient) GetLogs(
	pipelineName string,
	jobId string,
	data []string,
) *LogsIter {
	request := pps.GetLogsRequest{}
	resp := &LogsIter{}
	if len(pipelineName) > 0 {
		request.Pipeline = &pps.Pipeline{pipelineName}
	}
	if len(jobId) > 0 {
		request.Job = &pps.Job{jobId}
	}
	resp.logsClient, resp.err = c.PpsAPIClient.GetLogs(c.ctx(), &request)
	return resp
}

// CreatePipeline creates a new pipeline, pipelines are the main computation
// object in PPS they create a flow of data from a set of input Repos to an
// output Repo (which has the same name as the pipeline). Whenever new data is
// committed to one of the input repos the pipelines will create jobs to bring
// the output Repo up to data.
// image is the Docker image to run the jobs in.
// cmd is the command passed to the Docker run invocation.
// NOTE as with Docker cmd is not run inside a shell that means that things
// like wildcard globbing (*), pipes (|) and file redirects (> and >>) will not
// work. To get that behavior you should have your command be a shell of your
// choice and pass a shell script to stdin.
// stdin is a slice of lines that are sent to your command on stdin. Lines need
// not end in newline characters.
// parallelism is how many copies of your container should run in parallel. You
// may pass 0 for parallelism in which case PPS will set the parallelism based
// on availabe resources.
// inputs specifies a set of Repos that will be visible to the jobs during runtime.
// commits to these repos will cause the pipeline to create new jobs to process them.
// update indicates that you want to update an existing pipeline
func (c APIClient) CreatePipeline(
	name string,
	image string,
	cmd []string,
	stdin []string,
	parallelismSpec *pps.ParallelismSpec,
	inputs []*pps.PipelineInput,
	outputBranch string,
	update bool,
) error {
	_, err := c.PpsAPIClient.CreatePipeline(
		c.ctx(),
		&pps.CreatePipelineRequest{
			Pipeline: NewPipeline(name),
			Transform: &pps.Transform{
				Image: image,
				Cmd:   cmd,
				Stdin: stdin,
			},
			ParallelismSpec: parallelismSpec,
			Inputs:          inputs,
			OutputBranch:    outputBranch,
			Update:          update,
		},
	)
	return sanitizeErr(err)
}

// InspectPipeline returns info about a specific pipeline.
func (c APIClient) InspectPipeline(pipelineName string) (*pps.PipelineInfo, error) {
	pipelineInfo, err := c.PpsAPIClient.InspectPipeline(
		c.ctx(),
		&pps.InspectPipelineRequest{
			Pipeline: NewPipeline(pipelineName),
		},
	)
	return pipelineInfo, sanitizeErr(err)
}

// ListPipeline returns info about all pipelines.
func (c APIClient) ListPipeline() ([]*pps.PipelineInfo, error) {
	pipelineInfos, err := c.PpsAPIClient.ListPipeline(
		c.ctx(),
		&pps.ListPipelineRequest{},
	)
	if err != nil {
		return nil, sanitizeErr(err)
	}
	return pipelineInfos.PipelineInfo, nil
}

// DeletePipeline deletes a pipeline along with its output Repo.
func (c APIClient) DeletePipeline(name string) error {
	_, err := c.PpsAPIClient.DeletePipeline(
		c.ctx(),
		&pps.DeletePipelineRequest{
			Pipeline: NewPipeline(name),
		},
	)
	return sanitizeErr(err)
}

// StartPipeline restarts a stopped pipeline.
func (c APIClient) StartPipeline(name string) error {
	_, err := c.PpsAPIClient.StartPipeline(
		c.ctx(),
		&pps.StartPipelineRequest{
			Pipeline: NewPipeline(name),
		},
	)
	return sanitizeErr(err)
}

// StopPipeline prevents a pipeline from processing things, it can be restarted
// with StartPipeline.
func (c APIClient) StopPipeline(name string) error {
	_, err := c.PpsAPIClient.StopPipeline(
		c.ctx(),
		&pps.StopPipelineRequest{
			Pipeline: NewPipeline(name),
		},
	)
	return sanitizeErr(err)
}

// RerunPipeline reruns a pipeline over a given set of commits. Exclude and
// include are filters that either include or exclude the ancestors of the
// given commits.  A commit is considered the ancestor of itself. The behavior
// is the same as that of ListCommit.
func (c APIClient) RerunPipeline(name string, include []*pfs.Commit, exclude []*pfs.Commit) error {
	_, err := c.PpsAPIClient.RerunPipeline(
		c.ctx(),
		&pps.RerunPipelineRequest{
			Pipeline: NewPipeline(name),
			Include:  include,
			Exclude:  exclude,
		},
	)
	return sanitizeErr(err)
}
<commit_msg>Add comments to log iterator<commit_after>package client

import (
	"fmt"
	"io"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pps"
)

// NewJob creates a pps.Job.
func NewJob(jobID string) *pps.Job {
	return &pps.Job{ID: jobID}
}

const (
	// PPSEtcdPrefixEnv is the environment variable that specifies the etcd
	// prefix that PPS uses.
	PPSEtcdPrefixEnv = "PPS_ETCD_PREFIX"
	// PPSWorkerIPEnv is the environment variable that a worker can use to
	// see its own IP.  The IP address is made available through the
	// Kubernetes downward API.
	PPSWorkerIPEnv = "PPS_WORKER_IP"
	// PPSPipelineNameEnv is the env var that sets the name of the pipeline
	// that the workers are running.
	PPSPipelineNameEnv = "PPS_PIPELINE_NAME"
	// PPSJobIDEnv is the env var that sets the ID of the job that the
	// workers are running (if the workers belong to an orphan job, rather than a
	// pipeline).
	PPSJobIDEnv = "PPS_JOB_ID"
	// PPSInputPrefix is the prefix of the path where datums are downloaded
	// to.  A datum of an input named `XXX` is downloaded to `/pfs/XXX/`.
	PPSInputPrefix = "/pfs"
	// PPSOutputPath is the path where the user code is
	// expected to write its output to.
	PPSOutputPath = "/pfs/out"
	// PPSWorkerPort is the port that workers use for their gRPC server
	PPSWorkerPort = 30652
)

// NewJobInput creates a pps.JobInput.
func NewJobInput(repoName string, commitID string, glob string) *pps.JobInput {
	return &pps.JobInput{
		Commit: NewCommit(repoName, commitID),
		Glob:   glob,
	}
}

// NewPipeline creates a pps.Pipeline.
func NewPipeline(pipelineName string) *pps.Pipeline {
	return &pps.Pipeline{Name: pipelineName}
}

// NewPipelineInput creates a new pps.PipelineInput
func NewPipelineInput(repoName string, glob string) *pps.PipelineInput {
	return &pps.PipelineInput{
		Repo: NewRepo(repoName),
		Glob: glob,
	}
}

// CreateJob creates and runs a job in PPS.
// image is the Docker image to run the job in.
// cmd is the command passed to the Docker run invocation.
// NOTE as with Docker cmd is not run inside a shell that means that things
// like wildcard globbing (*), pipes (|) and file redirects (> and >>) will not
// work. To get that behavior you should have your command be a shell of your
// choice and pass a shell script to stdin.
// stdin is a slice of lines that are sent to your command on stdin. Lines need
// not end in newline characters.
// parallelism is how many copies of your container should run in parallel. You
// may pass 0 for parallelism in which case PPS will set the parallelism based
// on availabe resources.
// inputs specifies a set of Commits that will be visible to the job during runtime.
// parentJobID specifies the a job to use as a parent, it may be left empty in
// which case there is no parent job. If not left empty your job will use the
// parent Job's output commit as the parent of its output commit.
func (c APIClient) CreateJob(
	image string,
	cmd []string,
	stdin []string,
	parallelismSpec *pps.ParallelismSpec,
	inputs []*pps.JobInput,
	internalPort int32,
	externalPort int32,
) (*pps.Job, error) {
	var service *pps.Service
	if internalPort != 0 {
		service = &pps.Service{
			InternalPort: internalPort,
		}
	}
	if externalPort != 0 {
		if internalPort == 0 {
			return nil, fmt.Errorf("external port specified without internal port")
		}
		service.ExternalPort = externalPort
	}
	job, err := c.PpsAPIClient.CreateJob(
		c.ctx(),
		&pps.CreateJobRequest{
			Transform: &pps.Transform{
				Image: image,
				Cmd:   cmd,
				Stdin: stdin,
			},
			ParallelismSpec: parallelismSpec,
			Inputs:          inputs,
			Service:         service,
		},
	)
	return job, sanitizeErr(err)
}

// InspectJob returns info about a specific job.
// blockOutput will cause the call to block until the job has been assigned an output commit.
// blockState will cause the call to block until the job reaches a terminal state (failure or success).
func (c APIClient) InspectJob(jobID string, blockState bool) (*pps.JobInfo, error) {
	jobInfo, err := c.PpsAPIClient.InspectJob(
		c.ctx(),
		&pps.InspectJobRequest{
			Job:        NewJob(jobID),
			BlockState: blockState,
		})
	return jobInfo, sanitizeErr(err)
}

// ListJob returns info about all jobs.
// If pipelineName is non empty then only jobs that were started by the named pipeline will be returned
// If inputCommit is non-nil then only jobs which took the specific commits as inputs will be returned.
// The order of the inputCommits doesn't matter.
func (c APIClient) ListJob(pipelineName string, inputCommit []*pfs.Commit) ([]*pps.JobInfo, error) {
	var pipeline *pps.Pipeline
	if pipelineName != "" {
		pipeline = NewPipeline(pipelineName)
	}
	jobInfos, err := c.PpsAPIClient.ListJob(
		c.ctx(),
		&pps.ListJobRequest{
			Pipeline:    pipeline,
			InputCommit: inputCommit,
		})
	if err != nil {
		return nil, sanitizeErr(err)
	}
	return jobInfos.JobInfo, nil
}

// DeleteJob deletes a job along with its output Repo
func (c APIClient) DeleteJob(jobID string) error {
	_, err := c.PpsAPIClient.DeleteJob(
		c.ctx(),
		&pps.DeleteJobRequest{
			Job: NewJob(jobID),
		},
	)
	return sanitizeErr(err)
}

// LogsIter iterates through log messages returned from pps.GetLogs. Logs can
// be fetched with 'Next()'. The log message received can be examined with
// 'Message()', and any errors can be examined with 'Err()'.
type LogsIter struct {
	logsClient pps.API_GetLogsClient
	msg        *pps.LogMessage
	err        error
}

// Next retrieves the next relevant log message from pachd
func (l *LogsIter) Next() bool {
	if l.err != nil {
		l.msg = nil
		return false
	}
	l.msg, l.err = l.logsClient.Recv()
	if l.err != nil {
		return false
	}
	return true
}

// Message returns the most recently retrieve log message (as an annotated log
// line, in the form of a pps.LogMessage)
func (l *LogsIter) Message() *pps.LogMessage {
	return l.msg
}

// Err retrieves any errors encountered in the course of calling 'Next()'.
func (l *LogsIter) Err() error {
	if l.err == io.EOF {
		return nil
	}
	return l.err
}

// GetLogs gets logs from a job (logs includes stdout and stderr). 'pipelineName',
// 'jobID', and 'data', are all filters. To forego any filter, simply pass an
// empty value, though one of 'pipelineName' and 'jobID' must be set. Responses
// are written to 'messages'
func (c APIClient) GetLogs(
	pipelineName string,
	jobID string,
	data []string,
) *LogsIter {
	request := pps.GetLogsRequest{}
	resp := &LogsIter{}
	if len(pipelineName) > 0 {
		request.Pipeline = &pps.Pipeline{pipelineName}
	}
	if len(jobID) > 0 {
		request.Job = &pps.Job{jobID}
	}
	resp.logsClient, resp.err = c.PpsAPIClient.GetLogs(c.ctx(), &request)
	return resp
}

// CreatePipeline creates a new pipeline, pipelines are the main computation
// object in PPS they create a flow of data from a set of input Repos to an
// output Repo (which has the same name as the pipeline). Whenever new data is
// committed to one of the input repos the pipelines will create jobs to bring
// the output Repo up to data.
// image is the Docker image to run the jobs in.
// cmd is the command passed to the Docker run invocation.
// NOTE as with Docker cmd is not run inside a shell that means that things
// like wildcard globbing (*), pipes (|) and file redirects (> and >>) will not
// work. To get that behavior you should have your command be a shell of your
// choice and pass a shell script to stdin.
// stdin is a slice of lines that are sent to your command on stdin. Lines need
// not end in newline characters.
// parallelism is how many copies of your container should run in parallel. You
// may pass 0 for parallelism in which case PPS will set the parallelism based
// on availabe resources.
// inputs specifies a set of Repos that will be visible to the jobs during runtime.
// commits to these repos will cause the pipeline to create new jobs to process them.
// update indicates that you want to update an existing pipeline
func (c APIClient) CreatePipeline(
	name string,
	image string,
	cmd []string,
	stdin []string,
	parallelismSpec *pps.ParallelismSpec,
	inputs []*pps.PipelineInput,
	outputBranch string,
	update bool,
) error {
	_, err := c.PpsAPIClient.CreatePipeline(
		c.ctx(),
		&pps.CreatePipelineRequest{
			Pipeline: NewPipeline(name),
			Transform: &pps.Transform{
				Image: image,
				Cmd:   cmd,
				Stdin: stdin,
			},
			ParallelismSpec: parallelismSpec,
			Inputs:          inputs,
			OutputBranch:    outputBranch,
			Update:          update,
		},
	)
	return sanitizeErr(err)
}

// InspectPipeline returns info about a specific pipeline.
func (c APIClient) InspectPipeline(pipelineName string) (*pps.PipelineInfo, error) {
	pipelineInfo, err := c.PpsAPIClient.InspectPipeline(
		c.ctx(),
		&pps.InspectPipelineRequest{
			Pipeline: NewPipeline(pipelineName),
		},
	)
	return pipelineInfo, sanitizeErr(err)
}

// ListPipeline returns info about all pipelines.
func (c APIClient) ListPipeline() ([]*pps.PipelineInfo, error) {
	pipelineInfos, err := c.PpsAPIClient.ListPipeline(
		c.ctx(),
		&pps.ListPipelineRequest{},
	)
	if err != nil {
		return nil, sanitizeErr(err)
	}
	return pipelineInfos.PipelineInfo, nil
}

// DeletePipeline deletes a pipeline along with its output Repo.
func (c APIClient) DeletePipeline(name string) error {
	_, err := c.PpsAPIClient.DeletePipeline(
		c.ctx(),
		&pps.DeletePipelineRequest{
			Pipeline: NewPipeline(name),
		},
	)
	return sanitizeErr(err)
}

// StartPipeline restarts a stopped pipeline.
func (c APIClient) StartPipeline(name string) error {
	_, err := c.PpsAPIClient.StartPipeline(
		c.ctx(),
		&pps.StartPipelineRequest{
			Pipeline: NewPipeline(name),
		},
	)
	return sanitizeErr(err)
}

// StopPipeline prevents a pipeline from processing things, it can be restarted
// with StartPipeline.
func (c APIClient) StopPipeline(name string) error {
	_, err := c.PpsAPIClient.StopPipeline(
		c.ctx(),
		&pps.StopPipelineRequest{
			Pipeline: NewPipeline(name),
		},
	)
	return sanitizeErr(err)
}

// RerunPipeline reruns a pipeline over a given set of commits. Exclude and
// include are filters that either include or exclude the ancestors of the
// given commits.  A commit is considered the ancestor of itself. The behavior
// is the same as that of ListCommit.
func (c APIClient) RerunPipeline(name string, include []*pfs.Commit, exclude []*pfs.Commit) error {
	_, err := c.PpsAPIClient.RerunPipeline(
		c.ctx(),
		&pps.RerunPipelineRequest{
			Pipeline: NewPipeline(name),
			Include:  include,
			Exclude:  exclude,
		},
	)
	return sanitizeErr(err)
}
<|endoftext|>


<commit_before>package lba

import (
	"bytes"
	"errors"
	"fmt"
	"sync"

	"github.com/garyburd/redigo/redis"
	"github.com/zero-os/0-Disk"
	"github.com/zero-os/0-Disk/log"
)

// MetaRedisProvider is used by the LBA,
// to retreive a Redis Meta Connection
type MetaRedisProvider interface {
	MetaRedisConnection() (redis.Conn, error)
}

//NewLBA creates a new LBA
func NewLBA(vdiskID string, blockCount, cacheLimitInBytes int64, provider MetaRedisProvider) (lba *LBA, err error) {
	if provider == nil {
		return nil, errors.New("NewLBA requires a non-nil MetaRedisProvider")
	}

	muxCount := blockCount / NumberOfRecordsPerLBAShard
	if blockCount%NumberOfRecordsPerLBAShard > 0 {
		muxCount++
	}

	lba = &LBA{
		provider: provider,
		vdiskID:  vdiskID,
		shardMux: make([]sync.Mutex, muxCount),
	}

	lba.cache, err = newShardCache(cacheLimitInBytes, lba.onCacheEviction)

	return
}

// LBA implements the functionality to lookup block keys through the logical block index.
// The data is persisted to an external metadataserver in shards of n keys,
// where n = NumberOfRecordsPerLBAShard.
type LBA struct {
	cache *shardCache

	// One mutex per shard, allows us to only lock
	// on a per-shard basis. Even with 65k block, that's still only a ~500 element mutex array.
	// We stil need to lock on a per-shard basis,
	// as otherwise we might have a race condition where for example
	// 2 operations might create a new shard, and thus we would miss an operation.
	shardMux []sync.Mutex

	provider MetaRedisProvider
	vdiskID  string
}

//Set the content hash for a specific block.
// When a key is updated, the shard containing this blockindex is marked as dirty and will be
// stored in the external metadataserver when Flush is called.
func (lba *LBA) Set(blockIndex int64, h zerodisk.Hash) (err error) {
	//Fetch the appropriate shard
	shard, err := func(shardIndex int64) (shard *shard, err error) {
		lba.shardMux[shardIndex].Lock()
		defer lba.shardMux[shardIndex].Unlock()

		shard, err = lba.getShard(shardIndex)
		if err != nil {
			return
		}
		if shard == nil {
			shard = newShard()
			// store the new shard in the cache,
			// otherwise it will be forgotten...
			lba.cache.Add(shardIndex, shard)
		}

		return
	}(blockIndex / NumberOfRecordsPerLBAShard)

	if err != nil {
		return
	}

	//Update the hash
	hashIndex := blockIndex % NumberOfRecordsPerLBAShard
	shard.Set(hashIndex, h)

	return
}

//Delete the content hash for a specific block.
// When a key is updated, the shard containing this blockindex is marked as dirty and will be
// stored in the external metadaserver when Flush is called
// Deleting means actually that the nilhash will be set for this blockindex.
func (lba *LBA) Delete(blockIndex int64) (err error) {
	err = lba.Set(blockIndex, nil)
	return
}

//Get returns the hash for a block, nil if no hash registered
// If the shard containing this blockindex is not present, it is fetched from the external metadaserver
func (lba *LBA) Get(blockIndex int64) (h zerodisk.Hash, err error) {
	shard, err := func(shardIndex int64) (*shard, error) {
		lba.shardMux[shardIndex].Lock()
		defer lba.shardMux[shardIndex].Unlock()

		return lba.getShard(shardIndex)
	}(blockIndex / NumberOfRecordsPerLBAShard)

	if err != nil || shard == nil {
		return
	}

	// get the hash
	hashIndex := blockIndex % NumberOfRecordsPerLBAShard
	h = shard.Get(hashIndex)
	if h.Equals(zerodisk.NilHash) {
		h = nil
	}

	return
}

//Flush stores all dirty shards to the external metadaserver
func (lba *LBA) Flush() (err error) {
	err = lba.storeCacheInExternalStorage()
	return
}

func (lba *LBA) getShard(index int64) (shard *shard, err error) {
	shard, ok := lba.cache.Get(index)
	if !ok {
		shard, err = lba.getShardFromExternalStorage(index)
		if err != nil {
			return
		}

		if shard != nil {
			lba.cache.Add(index, shard)
		}
	}

	return
}

// in case a shard gets evicted from cache,
// this method will be called, and we'll serialize the shard immediately,
// unless it isn't dirty
func (lba *LBA) onCacheEviction(index int64, shard *shard) {
	if !shard.Dirty() {
		return
	}

	var err error

	// the given shard can be nil in case it was deleted by the user,
	// in that case we will remove the shard from the external storage as well
	// otherwise we serialize the shard before it gets thrown into the void
	if shard != nil {
		err = lba.storeShardInExternalStorage(index, shard)
	} else {
		err = lba.deleteShardFromExternalStorage(index)
	}

	if err != nil {
		log.Infof("error during eviction of shard %d: %s", index, err)
	}
}

func (lba *LBA) getShardFromExternalStorage(index int64) (shard *shard, err error) {
	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()
	reply, err := conn.Do("HGET", lba.vdiskID, index)
	if err != nil || reply == nil {
		return
	}

	shardBytes, err := redis.Bytes(reply, err)
	if err != nil {
		return
	}

	shard, err = shardFromBytes(shardBytes)
	return
}

func (lba *LBA) storeCacheInExternalStorage() (err error) {
	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()

	var cmdCount int64
	lba.cache.Serialize(func(index int64, bytes []byte) (err error) {
		if bytes != nil {
			err = conn.Send("HSET", lba.vdiskID, index, bytes)
		} else {
			err = conn.Send("HDEL", lba.vdiskID, index)
		}

		cmdCount++
		return
	})

	// Write all sets in output buffer to Redis at once
	err = conn.Flush()
	if err != nil {
		return
	}

	// read all responses
	var errors flushError
	for i := int64(0); i < cmdCount; i++ {
		_, err = conn.Receive()
		if err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		err = errors // return 1+ errors in case we received any
		return
	}

	// no need to evict, already serialized them
	evict := false
	// clear cache, as we serialized them all
	lba.cache.Clear(evict)

	// return with no errors, all good
	return
}

func (lba *LBA) storeShardInExternalStorage(index int64, shard *shard) (err error) {
	if !shard.Dirty() {
		log.Debugf(
			"LBA shard %d for %s isn't dirty, so nothing to store in external (meta) storage",
			index, lba.vdiskID)
		return // only store a dirty shard
	}

	var buffer bytes.Buffer
	if err = shard.Write(&buffer); err != nil {
		err = fmt.Errorf("couldn't serialize evicted shard %d: %s", index, err)
		return
	}

	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()

	_, err = conn.Do("HSET", lba.vdiskID, index, buffer.Bytes())
	if err != nil {
		shard.UnsetDirty()
	}

	return
}

func (lba *LBA) deleteShardFromExternalStorage(index int64) (err error) {
	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()

	_, err = conn.Do("HDEL", lba.vdiskID, index)

	return
}

// flushError is a collection of errors received,
// send by the server as a reply on the commands that got flushed to it
type flushError []error

// Error implements Error.Error
func (e flushError) Error() (s string) {
	s = fmt.Sprintf("flush failed because of %d commands: ", len(e))
	for _, err := range e {
		s += `"` + err.Error() + `";`
	}
	return
}
<commit_msg>Protect lba Get and Set with lock<commit_after>package lba

import (
	"bytes"
	"errors"
	"fmt"
	"sync"

	"github.com/garyburd/redigo/redis"
	"github.com/zero-os/0-Disk"
	"github.com/zero-os/0-Disk/log"
)

// MetaRedisProvider is used by the LBA,
// to retreive a Redis Meta Connection
type MetaRedisProvider interface {
	MetaRedisConnection() (redis.Conn, error)
}

//NewLBA creates a new LBA
func NewLBA(vdiskID string, blockCount, cacheLimitInBytes int64, provider MetaRedisProvider) (lba *LBA, err error) {
	if provider == nil {
		return nil, errors.New("NewLBA requires a non-nil MetaRedisProvider")
	}

	muxCount := blockCount / NumberOfRecordsPerLBAShard
	if blockCount%NumberOfRecordsPerLBAShard > 0 {
		muxCount++
	}

	lba = &LBA{
		provider: provider,
		vdiskID:  vdiskID,
		shardMux: make([]sync.RWMutex, muxCount),
	}

	lba.cache, err = newShardCache(cacheLimitInBytes, lba.onCacheEviction)

	return
}

// LBA implements the functionality to lookup block keys through the logical block index.
// The data is persisted to an external metadataserver in shards of n keys,
// where n = NumberOfRecordsPerLBAShard.
type LBA struct {
	cache *shardCache

	// One mutex per shard, allows us to only lock
	// on a per-shard basis. Even with 65k block, that's still only a ~500 element mutex array.
	// We stil need to lock on a per-shard basis,
	// as otherwise we might have a race condition where for example
	// 2 operations might create a new shard, and thus we would miss an operation.
	shardMux []sync.RWMutex

	provider MetaRedisProvider
	vdiskID  string
}

//Set the content hash for a specific block.
// When a key is updated, the shard containing this blockindex is marked as dirty and will be
// stored in the external metadataserver when Flush is called.
func (lba *LBA) Set(blockIndex int64, h zerodisk.Hash) (err error) {
	shardIndex := blockIndex / NumberOfRecordsPerLBAShard
	//Fetch the appropriate shard
	shard, err := func(shardIndex int64) (shard *shard, err error) {
		lba.shardMux[shardIndex].Lock()
		defer lba.shardMux[shardIndex].Unlock()

		shard, err = lba.getShard(shardIndex)
		if err != nil {
			return
		}
		if shard == nil {
			shard = newShard()
			// store the new shard in the cache,
			// otherwise it will be forgotten...
			lba.cache.Add(shardIndex, shard)
		}

		return
	}(shardIndex)

	if err != nil {
		return
	}

	//Update the hash
	hashIndex := blockIndex % NumberOfRecordsPerLBAShard

	lba.shardMux[shardIndex].Lock()
	defer lba.shardMux[shardIndex].Unlock()

	shard.Set(hashIndex, h)

	return
}

//Delete the content hash for a specific block.
// When a key is updated, the shard containing this blockindex is marked as dirty and will be
// stored in the external metadaserver when Flush is called
// Deleting means actually that the nilhash will be set for this blockindex.
func (lba *LBA) Delete(blockIndex int64) (err error) {
	err = lba.Set(blockIndex, nil)
	return
}

//Get returns the hash for a block, nil if no hash registered
// If the shard containing this blockindex is not present, it is fetched from the external metadaserver
func (lba *LBA) Get(blockIndex int64) (h zerodisk.Hash, err error) {
	shardIndex := blockIndex / NumberOfRecordsPerLBAShard

	shard, err := func(shardIndex int64) (*shard, error) {
		lba.shardMux[shardIndex].RLock()
		defer lba.shardMux[shardIndex].RUnlock()

		return lba.getShard(shardIndex)
	}(shardIndex)

	if err != nil || shard == nil {
		return
	}

	// get the hash
	hashIndex := blockIndex % NumberOfRecordsPerLBAShard

	lba.shardMux[shardIndex].RLock()
	defer lba.shardMux[shardIndex].RUnlock()

	h = shard.Get(hashIndex)
	if h.Equals(zerodisk.NilHash) {
		h = nil
	}

	return
}

//Flush stores all dirty shards to the external metadaserver
func (lba *LBA) Flush() (err error) {
	err = lba.storeCacheInExternalStorage()
	return
}

func (lba *LBA) getShard(index int64) (shard *shard, err error) {
	shard, ok := lba.cache.Get(index)
	if !ok {
		shard, err = lba.getShardFromExternalStorage(index)
		if err != nil {
			return
		}

		if shard != nil {
			lba.cache.Add(index, shard)
		}
	}

	return
}

// in case a shard gets evicted from cache,
// this method will be called, and we'll serialize the shard immediately,
// unless it isn't dirty
func (lba *LBA) onCacheEviction(index int64, shard *shard) {
	if !shard.Dirty() {
		return
	}

	var err error

	// the given shard can be nil in case it was deleted by the user,
	// in that case we will remove the shard from the external storage as well
	// otherwise we serialize the shard before it gets thrown into the void
	if shard != nil {
		err = lba.storeShardInExternalStorage(index, shard)
	} else {
		err = lba.deleteShardFromExternalStorage(index)
	}

	if err != nil {
		log.Infof("error during eviction of shard %d: %s", index, err)
	}
}

func (lba *LBA) getShardFromExternalStorage(index int64) (shard *shard, err error) {
	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()
	reply, err := conn.Do("HGET", lba.vdiskID, index)
	if err != nil || reply == nil {
		return
	}

	shardBytes, err := redis.Bytes(reply, err)
	if err != nil {
		return
	}

	shard, err = shardFromBytes(shardBytes)
	return
}

func (lba *LBA) storeCacheInExternalStorage() (err error) {
	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()

	var cmdCount int64
	lba.cache.Serialize(func(index int64, bytes []byte) (err error) {
		if bytes != nil {
			err = conn.Send("HSET", lba.vdiskID, index, bytes)
		} else {
			err = conn.Send("HDEL", lba.vdiskID, index)
		}

		cmdCount++
		return
	})

	// Write all sets in output buffer to Redis at once
	err = conn.Flush()
	if err != nil {
		return
	}

	// read all responses
	var errors flushError
	for i := int64(0); i < cmdCount; i++ {
		_, err = conn.Receive()
		if err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		err = errors // return 1+ errors in case we received any
		return
	}

	// no need to evict, already serialized them
	evict := false
	// clear cache, as we serialized them all
	lba.cache.Clear(evict)

	// return with no errors, all good
	return
}

func (lba *LBA) storeShardInExternalStorage(index int64, shard *shard) (err error) {
	if !shard.Dirty() {
		log.Debugf(
			"LBA shard %d for %s isn't dirty, so nothing to store in external (meta) storage",
			index, lba.vdiskID)
		return // only store a dirty shard
	}

	var buffer bytes.Buffer
	if err = shard.Write(&buffer); err != nil {
		err = fmt.Errorf("couldn't serialize evicted shard %d: %s", index, err)
		return
	}

	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()

	_, err = conn.Do("HSET", lba.vdiskID, index, buffer.Bytes())
	if err != nil {
		shard.UnsetDirty()
	}

	return
}

func (lba *LBA) deleteShardFromExternalStorage(index int64) (err error) {
	conn, err := lba.provider.MetaRedisConnection()
	if err != nil {
		return
	}
	defer conn.Close()

	_, err = conn.Do("HDEL", lba.vdiskID, index)

	return
}

// flushError is a collection of errors received,
// send by the server as a reply on the commands that got flushed to it
type flushError []error

// Error implements Error.Error
func (e flushError) Error() (s string) {
	s = fmt.Sprintf("flush failed because of %d commands: ", len(e))
	for _, err := range e {
		s += `"` + err.Error() + `";`
	}
	return
}
<|endoftext|>


<commit_before>package dallimin_test

import (
	"encoding/json"
	"io/ioutil"
	"strconv"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/subosito/dallimin"
)

type Result struct {
	Server string `json:"server"`
	Key    string `json:"key"`
	Weight int    `json:"weight"`
}

type Fixture struct {
	Results []Result
	Servers []string
	Keys    []string
}

func loadFixture(fname string) Fixture {
	file, err := ioutil.ReadFile(fname)
	panicErr(err)

	var fixture Fixture

	err = json.Unmarshal(file, &fixture)
	panicErr(err)

	return fixture
}

func panicErr(err error) {
	if err != nil {
		panic(err)
	}
}

func TestPickServer(t *testing.T) {
	f := loadFixture("fixtures/keys.json")
	h, _ := dallimin.New(f.Servers, dallimin.Option{})

	for _, data := range f.Results {
		addr, err := h.PickServer(data.Key)
		server := strings.Split(data.Server, ":")

		assert.Nil(t, err)
		assert.Equal(t, "127.0.0.1:"+server[1], addr.String())
	}
}

func TestPickServer_withWeights(t *testing.T) {
	f := loadFixture("fixtures/keys-with-weights.json")
	s := map[string]int{}

	for _, server := range f.Servers {
		w := strings.Split(server, ":")[2]
		v := strings.TrimSuffix(server, ":"+w)

		n, _ := strconv.Atoi(w)

		s[v] = n
	}

	h, _ := dallimin.NewWithWeights(s, dallimin.Option{})

	for _, data := range f.Results {
		addr, err := h.PickServer(data.Key)
		server := strings.Split(data.Server, ":")

		assert.Nil(t, err)
		assert.Equal(t, "127.0.0.1:"+server[1], addr.String())
	}
}

func TestPickServer_singleServer(t *testing.T) {
	s := []string{"127.0.0.1:11211"}
	h, _ := dallimin.New(s, dallimin.Option{})

	addr, err := h.PickServer("api:foo")

	assert.Nil(t, err)
	assert.Equal(t, "127.0.0.1:11211", addr.String())
}

func TestPickServer_noServer(t *testing.T) {
	s := []string{}
	h, _ := dallimin.New(s, dallimin.Option{})

	addr, err := h.PickServer("api:foo")

	assert.Equal(t, err, dallimin.ErrNoServers)
	assert.Nil(t, addr)
}

func TestPickServer_whenNoServerAlive(t *testing.T) {
	s := []string{
		"127.0.0.1:12345",
		"127.0.0.1:12346",
	}

	h, _ := dallimin.New(s, dallimin.Option{CheckAlive: true})

	addr, err := h.PickServer("api:foo")
	assert.Equal(t, err, dallimin.ErrNoServers)
	assert.Nil(t, addr)
}
<commit_msg>added check alive and failover tests<commit_after>package dallimin_test

import (
	"encoding/json"
	"io/ioutil"
	"strconv"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/subosito/dallimin"
)

type Result struct {
	Server string `json:"server"`
	Key    string `json:"key"`
	Weight int    `json:"weight"`
}

type Fixture struct {
	Results []Result
	Servers []string
	Keys    []string
}

func loadFixture(fname string) Fixture {
	file, err := ioutil.ReadFile(fname)
	panicErr(err)

	var fixture Fixture

	err = json.Unmarshal(file, &fixture)
	panicErr(err)

	return fixture
}

func panicErr(err error) {
	if err != nil {
		panic(err)
	}
}

func TestPickServer(t *testing.T) {
	f := loadFixture("fixtures/keys.json")
	h, _ := dallimin.New(f.Servers, dallimin.Option{})

	for _, data := range f.Results {
		addr, err := h.PickServer(data.Key)
		server := strings.Split(data.Server, ":")

		assert.Nil(t, err)
		assert.Equal(t, "127.0.0.1:"+server[1], addr.String())
	}
}

func TestPickServer_withWeights(t *testing.T) {
	f := loadFixture("fixtures/keys-with-weights.json")
	s := map[string]int{}

	for _, server := range f.Servers {
		w := strings.Split(server, ":")[2]
		v := strings.TrimSuffix(server, ":"+w)

		n, _ := strconv.Atoi(w)

		s[v] = n
	}

	h, _ := dallimin.NewWithWeights(s, dallimin.Option{})

	for _, data := range f.Results {
		addr, err := h.PickServer(data.Key)
		server := strings.Split(data.Server, ":")

		assert.Nil(t, err)
		assert.Equal(t, "127.0.0.1:"+server[1], addr.String())
	}
}

func TestPickServer_singleServer(t *testing.T) {
	s := []string{"127.0.0.1:11211"}
	h, _ := dallimin.New(s, dallimin.Option{})

	addr, err := h.PickServer("api:foo")

	assert.Nil(t, err)
	assert.Equal(t, "127.0.0.1:11211", addr.String())
}

func TestPickServer_noServer(t *testing.T) {
	s := []string{}
	h, _ := dallimin.New(s, dallimin.Option{})

	addr, err := h.PickServer("api:foo")

	assert.Equal(t, err, dallimin.ErrNoServers)
	assert.Nil(t, addr)
}

func TestPickServer_whenNoServerAlive(t *testing.T) {
	s := []string{
		"127.0.0.1:12345",
		"127.0.0.1:12346",
	}

	h, _ := dallimin.New(s, dallimin.Option{CheckAlive: true})

	addr, err := h.PickServer("api:foo")
	assert.Equal(t, err, dallimin.ErrNoServers)
	assert.Nil(t, addr)
}

func TestPickServer_whenAtLeastOneAlive(t *testing.T) {
	s := []string{
		"127.0.0.1:12345",
		"127.0.0.1:11211",
	}

	h, _ := dallimin.New(s, dallimin.Option{CheckAlive: true})

	addr, err := h.PickServer("api:foo")
	assert.Nil(t, err)
	assert.Equal(t, addr.String(), "127.0.0.1:11211")
}

func TestPickServer_checkAliveAndfailover(t *testing.T) {
	s := []string{
		"127.0.0.1:12346",
		"127.0.0.1:11210",
	}

	h, _ := dallimin.New(s, dallimin.Option{CheckAlive: true, Failover: true})

	addr, err := h.PickServer("api:foo")
	assert.Nil(t, err)
	assert.Equal(t, addr.String(), "127.0.0.1:11210")
}
<|endoftext|>