<commit_before>package main

import (
	"flag"
	"log"
	"net/http"
	_ "net/http/pprof"
	"sync"
	"time"

	consul_api "github.com/hashicorp/consul/api"
	consul "github.com/hashicorp/consul/consul/structs"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	namespace = "consul"
)

var (
	serviceLabelNames = []string{"service", "node"}
	memberLabelNames  = []string{"member"}
)

// Exporter collects Consul stats from the given server and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI   string
	mutex sync.RWMutex

	up, clusterServers                     prometheus.Gauge
	nodeCount, serviceCount                prometheus.Counter
	serviceNodesTotal, serviceNodesHealthy *prometheus.GaugeVec
	client                                 *consul_api.Client
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string, consulLocks string, timeout time.Duration) *Exporter {
	// Set up our Consul client connection.
	consul_client, _ := consul_api.NewClient(&consul_api.Config{
		Address: uri,
	})

	// Init our exporter.
	return &Exporter{
		URI: uri,
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "up",
			Help:      "Was the last query of Consul successful.",
		}),

		clusterServers: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "raft_peers",
			Help:      "How many peers (servers) are in the Raft cluster.",
		}),

		nodeCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "serf_lan_members",
			Help:      "How many members are in the cluster.",
		}),

		serviceCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "catalog_services",
			Help:      "How many services are in the cluster.",
		}),

		serviceNodesTotal: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_nodes",
				Help:      "Number of nodes currently registered for this service",
			},
			[]string{"service"},
		),

		serviceNodesHealthy: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_node_healthy",
				Help:      "Is this service healthy on this node?",
			},
			[]string{"service", "node"},
		),

		client: consul_client,
	}
}

// Describe describes all the metrics ever exported by the Consul exporter. It
// implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.up.Desc()
	ch <- e.nodeCount.Desc()
	ch <- e.serviceCount.Desc()
	ch <- e.clusterServers.Desc()

	e.serviceNodesTotal.Describe(ch)
	e.serviceNodesHealthy.Describe(ch)
}

// Collect fetches the stats from configured Consul location and delivers them
// as Prometheus metrics. It implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	services := make(chan []*consul_api.ServiceEntry)

	go e.queryClient(services)

	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()

	// Reset metrics.
	e.serviceNodesTotal.Reset()
	e.serviceNodesHealthy.Reset()

	e.setMetrics(services)

	ch <- e.up
	ch <- e.clusterServers
	ch <- e.nodeCount
	ch <- e.serviceCount

	e.serviceNodesTotal.Collect(ch)
	e.serviceNodesHealthy.Collect(ch)
}

func (e *Exporter) queryClient(services chan<- []*consul_api.ServiceEntry) {
	defer close(services)

	// How many peers are in the Consul cluster?
	peers, err := e.client.Status().Peers()

	if err != nil {
		e.up.Set(0)
		log.Printf("Query error is %v", err)
		return
	}

	// We'll use peers to decide that we're up.
	e.up.Set(1)
	e.clusterServers.Set(float64(len(peers)))

	// How many nodes are registered?
	nodes, _, err := e.client.Catalog().Nodes(&consul_api.QueryOptions{})

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
	} else {
		e.nodeCount.Set(float64(len(nodes)))
	}

	// Query for the full list of services.
	serviceNames, _, err := e.client.Catalog().Services(&consul_api.QueryOptions{})
	e.serviceCount.Set(float64(len(serviceNames)))

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
		return
	}

	e.serviceCount.Set(float64(len(serviceNames)))

	for s := range serviceNames {
		s_entries, _, err := e.client.Health().Service(s, "", false, &consul_api.QueryOptions{})

		if err != nil {
			log.Printf("Failed to query service health: %v", err)
			continue
		}

		services <- s_entries
	}
}

func (e *Exporter) setMetrics(services <-chan []*consul_api.ServiceEntry) {
	// Each service will be an array of ServiceEntry structs.
	for service := range services {
		if len(service) == 0 {
			// Not sure this should ever happen, but catch it just in case...
			continue
		}

		// We should have one ServiceEntry per node, so use that for total nodes.
		e.serviceNodesTotal.WithLabelValues(service[0].Service.Service).Set(float64(len(service)))

		for _, entry := range service {
			// We have a Node, a Service, and one or more Checks. Our
			// service-node combo is passing if all checks have a `status`
			// of "passing."

			passing := 1

			for _, hc := range entry.Checks {
				if hc.Status != consul.HealthPassing {
					passing = 0
					break
				}
			}

			log.Printf("%v/%v status is %v", entry.Service.Service, entry.Node.Node, passing)

			e.serviceNodesHealthy.WithLabelValues(entry.Service.Service, entry.Node.Node).Set(float64(passing))
		}
	}
}

func main() {
	var (
		listenAddress = flag.String("web.listen-address", ":9107", "Address to listen on for web interface and telemetry.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
		consulServer  = flag.String("consul.server", "localhost:8500", "HTTP API address of a Consul server or agent.")
		//consulLocks   = flag.String("consul.locks", "", "If specified, keys to check for session locks. Comma-seperated list of keys.")
		//consulTimeout = flag.Duration("consul.timeout", 5*time.Second, "Timeout for trying to get stats from Consul.")
	)
	flag.Parse()

	exporter := NewExporter(*consulServer, "", 5*time.Second)
	prometheus.MustRegister(exporter)

	log.Printf("Starting Server: %s", *listenAddress)
	http.Handle(*metricsPath, prometheus.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
             <title>Consul Exporter</title>

Consul Exporter

<a href='` + *metricsPath + `'>Metrics

             `))
	})
	log.Fatal(http.ListenAndServe(*listenAddress, nil))
}
<commit_msg>Remove unused flags instead of commenting them out<commit_after>package main

import (
	"flag"
	"log"
	"net/http"
	_ "net/http/pprof"
	"sync"

	consul_api "github.com/hashicorp/consul/api"
	consul "github.com/hashicorp/consul/consul/structs"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	namespace = "consul"
)

var (
	serviceLabelNames = []string{"service", "node"}
	memberLabelNames  = []string{"member"}
)

// Exporter collects Consul stats from the given server and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI   string
	mutex sync.RWMutex

	up, clusterServers                     prometheus.Gauge
	nodeCount, serviceCount                prometheus.Counter
	serviceNodesTotal, serviceNodesHealthy *prometheus.GaugeVec
	client                                 *consul_api.Client
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string) *Exporter {
	// Set up our Consul client connection.
	consul_client, _ := consul_api.NewClient(&consul_api.Config{
		Address: uri,
	})

	// Init our exporter.
	return &Exporter{
		URI: uri,
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "up",
			Help:      "Was the last query of Consul successful.",
		}),

		clusterServers: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "raft_peers",
			Help:      "How many peers (servers) are in the Raft cluster.",
		}),

		nodeCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "serf_lan_members",
			Help:      "How many members are in the cluster.",
		}),

		serviceCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "catalog_services",
			Help:      "How many services are in the cluster.",
		}),

		serviceNodesTotal: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_nodes",
				Help:      "Number of nodes currently registered for this service",
			},
			[]string{"service"},
		),

		serviceNodesHealthy: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_node_healthy",
				Help:      "Is this service healthy on this node?",
			},
			[]string{"service", "node"},
		),

		client: consul_client,
	}
}

// Describe describes all the metrics ever exported by the Consul exporter. It
// implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.up.Desc()
	ch <- e.nodeCount.Desc()
	ch <- e.serviceCount.Desc()
	ch <- e.clusterServers.Desc()

	e.serviceNodesTotal.Describe(ch)
	e.serviceNodesHealthy.Describe(ch)
}

// Collect fetches the stats from configured Consul location and delivers them
// as Prometheus metrics. It implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	services := make(chan []*consul_api.ServiceEntry)

	go e.queryClient(services)

	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()

	// Reset metrics.
	e.serviceNodesTotal.Reset()
	e.serviceNodesHealthy.Reset()

	e.setMetrics(services)

	ch <- e.up
	ch <- e.clusterServers
	ch <- e.nodeCount
	ch <- e.serviceCount

	e.serviceNodesTotal.Collect(ch)
	e.serviceNodesHealthy.Collect(ch)
}

func (e *Exporter) queryClient(services chan<- []*consul_api.ServiceEntry) {
	defer close(services)

	// How many peers are in the Consul cluster?
	peers, err := e.client.Status().Peers()

	if err != nil {
		e.up.Set(0)
		log.Printf("Query error is %v", err)
		return
	}

	// We'll use peers to decide that we're up.
	e.up.Set(1)
	e.clusterServers.Set(float64(len(peers)))

	// How many nodes are registered?
	nodes, _, err := e.client.Catalog().Nodes(&consul_api.QueryOptions{})

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
	} else {
		e.nodeCount.Set(float64(len(nodes)))
	}

	// Query for the full list of services.
	serviceNames, _, err := e.client.Catalog().Services(&consul_api.QueryOptions{})
	e.serviceCount.Set(float64(len(serviceNames)))

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
		return
	}

	e.serviceCount.Set(float64(len(serviceNames)))

	for s := range serviceNames {
		s_entries, _, err := e.client.Health().Service(s, "", false, &consul_api.QueryOptions{})

		if err != nil {
			log.Printf("Failed to query service health: %v", err)
			continue
		}

		services <- s_entries
	}
}

func (e *Exporter) setMetrics(services <-chan []*consul_api.ServiceEntry) {
	// Each service will be an array of ServiceEntry structs.
	for service := range services {
		if len(service) == 0 {
			// Not sure this should ever happen, but catch it just in case...
			continue
		}

		// We should have one ServiceEntry per node, so use that for total nodes.
		e.serviceNodesTotal.WithLabelValues(service[0].Service.Service).Set(float64(len(service)))

		for _, entry := range service {
			// We have a Node, a Service, and one or more Checks. Our
			// service-node combo is passing if all checks have a `status`
			// of "passing."

			passing := 1

			for _, hc := range entry.Checks {
				if hc.Status != consul.HealthPassing {
					passing = 0
					break
				}
			}

			log.Printf("%v/%v status is %v", entry.Service.Service, entry.Node.Node, passing)

			e.serviceNodesHealthy.WithLabelValues(entry.Service.Service, entry.Node.Node).Set(float64(passing))
		}
	}
}

func main() {
	var (
		listenAddress = flag.String("web.listen-address", ":9107", "Address to listen on for web interface and telemetry.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
		consulServer  = flag.String("consul.server", "localhost:8500", "HTTP API address of a Consul server or agent.")
	)
	flag.Parse()

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

	log.Printf("Starting Server: %s", *listenAddress)
	http.Handle(*metricsPath, prometheus.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
             <title>Consul Exporter</title>

Consul Exporter

<a href='` + *metricsPath + `'>Metrics

             `))
	})
	log.Fatal(http.ListenAndServe(*listenAddress, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"strings"
)

func ExtractColumn(logContents string, column uint8) string {
	var buffer bytes.Buffer
	lines := strings.Split(logContents, "\n")

	for i := range lines {
        if lines[i] == "" {
            continue
        }

		cols := strings.Split(lines[i], " ")

		switch column {
		case 0:
			buffer.WriteString(cols[0])
			buffer.WriteString(" ")
			buffer.WriteString(cols[1])
		case 1:
			buffer.WriteString(cols[2])
		case 2:
			buffer.WriteString(cols[3])

			for j := 4; j < len(cols); j++ {
				buffer.WriteString(" ")
				buffer.WriteString(cols[j])
			}
		}

		buffer.WriteString("\n")
	}

	return buffer.String()
}

func main() {
	var (
		column uint8
		buffer bytes.Buffer
	)

	_, err := fmt.Scan(&column)

	if err != nil {
		return
	}

	scanner := bufio.NewScanner(os.Stdin)

	for scanner.Scan() {
		buffer.WriteString(scanner.Text())
		buffer.WriteString("\n")
	}

	fmt.Print(ExtractColumn(buffer.String(), column))
}
<commit_msg>optimize: replaced string splitting with use of Scanner for lines: no need to store all lines in an array<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"strings"
)

func ExtractColumn(logContents string, column uint8) string {
	var buffer bytes.Buffer
    scanner := bufio.NewScanner(strings.NewReader(logContents))

    for scanner.Scan() {
        line := scanner.Text()

        if line == "" {
            continue
        }

		cols := strings.Split(line, " ")

		switch column {
		case 0:
			buffer.WriteString(cols[0])
			buffer.WriteString(" ")
			buffer.WriteString(cols[1])
		case 1:
			buffer.WriteString(cols[2])
		case 2:
			buffer.WriteString(cols[3])

			for j := 4; j < len(cols); j++ {
				buffer.WriteString(" ")
				buffer.WriteString(cols[j])
			}
		}

		buffer.WriteString("\n")
	}

	return buffer.String()
}

func main() {
	var (
		column uint8
		buffer bytes.Buffer
	)

	_, err := fmt.Scan(&column)

	if err != nil {
		return
	}

	scanner := bufio.NewScanner(os.Stdin)

	for scanner.Scan() {
		buffer.WriteString(scanner.Text())
		buffer.WriteString("\n")
	}

	fmt.Print(ExtractColumn(buffer.String(), column))
}
<|endoftext|>


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

// Package pools provides functionality to manage and reuse resources
// like connections.
package pools

import (
	"errors"
	"fmt"
	"time"

	"github.com/youtube/vitess/go/sync2"
)

var (
	CLOSED_ERR  = errors.New("resource pool is closed")
	TIMEOUT_ERR = errors.New("resource pool timed out")
)

// Factory is a function that can be used to create a resource.
type Factory func() (Resource, error)

// Every resource needs to suport the Resource interface.
// Thread synchronization between Close() and IsClosed()
// is the responsibility the caller.
type Resource interface {
	Close()
}

// ResourcePool allows you to use a pool of resources.
type ResourcePool struct {
	resources   chan resourceWrapper
	factory     Factory
	capacity    sync2.AtomicInt64
	idleTimeout sync2.AtomicDuration

	// stats
	waitCount sync2.AtomicInt64
	waitTime  sync2.AtomicDuration
}

type resourceWrapper struct {
	resource Resource
	timeUsed time.Time
}

// NewResourcePool creates a new ResourcePool pool.
// capacity is the initial capacity of the pool.
// maxCap is the maximum capacity.
// If a resource is unused beyond idleTimeout, it's discarded.
// An idleTimeout of 0 means that there is no timeout.
func NewResourcePool(factory Factory, capacity, maxCap int, idleTimeout time.Duration) *ResourcePool {
	if capacity <= 0 || maxCap <= 0 || capacity > maxCap {
		panic(errors.New("invalid/out of range capacity"))
	}
	rp := &ResourcePool{
		resources:   make(chan resourceWrapper, maxCap),
		factory:     factory,
		capacity:    sync2.AtomicInt64(capacity),
		idleTimeout: sync2.AtomicDuration(idleTimeout),
	}
	for i := 0; i < capacity; i++ {
		rp.resources <- resourceWrapper{}
	}
	return rp
}

// Close empties the pool calling Close on all its resources.
// You can call Close while there are outstanding resources.
// It waits for all resources to be returned (Put).
// After a Close, Get and TryGet are not allowed.
func (rp *ResourcePool) Close() {
	_ = rp.SetCapacity(0)
}

func (rp *ResourcePool) IsClosed() (closed bool) {
	return rp.capacity.Get() == 0
}

// Get will return the next available resource. If capacity
// has not been reached, it will create a new one using the factory. Otherwise,
// it will wait till the next resource becomes available or a timeout.
// A timeout of 0 is an indefinite wait.
func (rp *ResourcePool) Get(timeout time.Duration) (resource Resource, err error) {
	return rp.get(true, timeout)
}

// TryGet will return the next available resource. If none is available, and capacity
// has not been reached, it will create a new one using the factory. Otherwise,
// it will return nil with no error.
func (rp *ResourcePool) TryGet() (resource Resource, err error) {
	return rp.get(false, 0)
}

func (rp *ResourcePool) get(wait bool, timeout time.Duration) (resource Resource, err error) {
	// Fetch
	var wrapper resourceWrapper
	var ok bool
	select {
	case wrapper, ok = <-rp.resources:
	default:
		if !wait {
			return nil, nil
		}
		startTime := time.Now()
		if timeout == 0 {
			wrapper, ok = <-rp.resources
		} else {
			tmr := time.NewTimer(timeout)
			defer tmr.Stop()
			select {
			case wrapper, ok = <-rp.resources:
			case <-tmr.C:
				return nil, TIMEOUT_ERR
			}
		}
		rp.recordWait(startTime)
	}
	if !ok {
		return nil, CLOSED_ERR
	}

	// Unwrap
	idleTimeout := rp.idleTimeout.Get()
	if wrapper.resource != nil && idleTimeout > 0 && wrapper.timeUsed.Add(idleTimeout).Sub(time.Now()) < 0 {
		wrapper.resource.Close()
		wrapper.resource = nil
	}
	if wrapper.resource == nil {
		wrapper.resource, err = rp.factory()
		if err != nil {
			rp.resources <- resourceWrapper{}
		}
	}
	return wrapper.resource, err
}

// Put will return a resource to the pool. For every successful Get,
// a corresponding Put is required. If you no longer need a resource,
// you will need to call Put(nil) instead of returning the closed resource.
// The will eventually cause a new resource to be created in its place.
func (rp *ResourcePool) Put(resource Resource) {
	var wrapper resourceWrapper
	if resource != nil {
		wrapper = resourceWrapper{resource, time.Now()}
	}
	select {
	case rp.resources <- wrapper:
	default:
		panic(errors.New("attempt to Put into a full ResourcePool"))
	}
}

// SetCapacity changes the capacity of the pool.
// You can use it to shrink or expand, but not beyond
// the max capacity. If the change requires the pool
// to be shrunk, SetCapacity waits till the necessary
// number of resources are returned to the pool.
// A SetCapacity of 0 is equivalent to closing the ResourcePool.
func (rp *ResourcePool) SetCapacity(capacity int) error {
	if capacity < 0 || capacity > cap(rp.resources) {
		return fmt.Errorf("capacity %d is out of range", capacity)
	}

	// Atomically swap new capacity with old, but only
	// if old capacity is non-zero.
	var oldcap int
	for {
		oldcap = int(rp.capacity.Get())
		if oldcap == 0 {
			return CLOSED_ERR
		}
		if oldcap == capacity {
			return nil
		}
		if rp.capacity.CompareAndSwap(int64(oldcap), int64(capacity)) {
			break
		}
	}

	if capacity < oldcap {
		for i := 0; i < oldcap-capacity; i++ {
			wrapper := <-rp.resources
			if wrapper.resource != nil {
				wrapper.resource.Close()
			}
		}
	} else {
		for i := 0; i < capacity-oldcap; i++ {
			rp.resources <- resourceWrapper{}
		}
	}
	if capacity == 0 {
		close(rp.resources)
	}
	return nil
}

func (rp *ResourcePool) recordWait(start time.Time) {
	rp.waitCount.Add(1)
	rp.waitTime.Add(time.Now().Sub(start))
}

func (rp *ResourcePool) SetIdleTimeout(idleTimeout time.Duration) {
	rp.idleTimeout.Set(idleTimeout)
}

func (rp *ResourcePool) StatsJSON() string {
	c, a, mx, wc, wt, it := rp.Stats()
	return fmt.Sprintf(`{"Capacity": %v, "Available": %v, "MaxCapacity": %v, "WaitCount": %v, "WaitTime": %v, "IdleTimeout": %v}`, c, a, mx, wc, int64(wt), int64(it))
}

func (rp *ResourcePool) Stats() (capacity, available, maxCap, waitCount int64, waitTime, idleTimeout time.Duration) {
	return rp.Capacity(), rp.Available(), rp.MaxCap(), rp.WaitCount(), rp.WaitTime(), rp.IdleTimeout()
}

func (rp *ResourcePool) Capacity() int64 {
	return rp.capacity.Get()
}

func (rp *ResourcePool) Available() int64 {
	return int64(len(rp.resources))
}

func (rp *ResourcePool) MaxCap() int64 {
	return int64(cap(rp.resources))
}

func (rp *ResourcePool) WaitCount() int64 {
	return rp.waitCount.Get()
}

func (rp *ResourcePool) WaitTime() time.Duration {
	return rp.waitTime.Get()
}

func (rp *ResourcePool) IdleTimeout() time.Duration {
	return rp.idleTimeout.Get()
}
<commit_msg>Grammar/typo correction<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package pools provides functionality to manage and reuse resources
// like connections.
package pools

import (
	"errors"
	"fmt"
	"time"

	"github.com/youtube/vitess/go/sync2"
)

var (
	CLOSED_ERR  = errors.New("resource pool is closed")
	TIMEOUT_ERR = errors.New("resource pool timed out")
)

// Factory is a function that can be used to create a resource.
type Factory func() (Resource, error)

// Every resource needs to suport the Resource interface.
// Thread synchronization between Close() and IsClosed()
// is the responsibility of the caller.
type Resource interface {
	Close()
}

// ResourcePool allows you to use a pool of resources.
type ResourcePool struct {
	resources   chan resourceWrapper
	factory     Factory
	capacity    sync2.AtomicInt64
	idleTimeout sync2.AtomicDuration

	// stats
	waitCount sync2.AtomicInt64
	waitTime  sync2.AtomicDuration
}

type resourceWrapper struct {
	resource Resource
	timeUsed time.Time
}

// NewResourcePool creates a new ResourcePool pool.
// capacity is the initial capacity of the pool.
// maxCap is the maximum capacity.
// If a resource is unused beyond idleTimeout, it's discarded.
// An idleTimeout of 0 means that there is no timeout.
func NewResourcePool(factory Factory, capacity, maxCap int, idleTimeout time.Duration) *ResourcePool {
	if capacity <= 0 || maxCap <= 0 || capacity > maxCap {
		panic(errors.New("invalid/out of range capacity"))
	}
	rp := &ResourcePool{
		resources:   make(chan resourceWrapper, maxCap),
		factory:     factory,
		capacity:    sync2.AtomicInt64(capacity),
		idleTimeout: sync2.AtomicDuration(idleTimeout),
	}
	for i := 0; i < capacity; i++ {
		rp.resources <- resourceWrapper{}
	}
	return rp
}

// Close empties the pool calling Close on all its resources.
// You can call Close while there are outstanding resources.
// It waits for all resources to be returned (Put).
// After a Close, Get and TryGet are not allowed.
func (rp *ResourcePool) Close() {
	_ = rp.SetCapacity(0)
}

func (rp *ResourcePool) IsClosed() (closed bool) {
	return rp.capacity.Get() == 0
}

// Get will return the next available resource. If capacity
// has not been reached, it will create a new one using the factory. Otherwise,
// it will wait till the next resource becomes available or a timeout.
// A timeout of 0 is an indefinite wait.
func (rp *ResourcePool) Get(timeout time.Duration) (resource Resource, err error) {
	return rp.get(true, timeout)
}

// TryGet will return the next available resource. If none is available, and capacity
// has not been reached, it will create a new one using the factory. Otherwise,
// it will return nil with no error.
func (rp *ResourcePool) TryGet() (resource Resource, err error) {
	return rp.get(false, 0)
}

func (rp *ResourcePool) get(wait bool, timeout time.Duration) (resource Resource, err error) {
	// Fetch
	var wrapper resourceWrapper
	var ok bool
	select {
	case wrapper, ok = <-rp.resources:
	default:
		if !wait {
			return nil, nil
		}
		startTime := time.Now()
		if timeout == 0 {
			wrapper, ok = <-rp.resources
		} else {
			tmr := time.NewTimer(timeout)
			defer tmr.Stop()
			select {
			case wrapper, ok = <-rp.resources:
			case <-tmr.C:
				return nil, TIMEOUT_ERR
			}
		}
		rp.recordWait(startTime)
	}
	if !ok {
		return nil, CLOSED_ERR
	}

	// Unwrap
	idleTimeout := rp.idleTimeout.Get()
	if wrapper.resource != nil && idleTimeout > 0 && wrapper.timeUsed.Add(idleTimeout).Sub(time.Now()) < 0 {
		wrapper.resource.Close()
		wrapper.resource = nil
	}
	if wrapper.resource == nil {
		wrapper.resource, err = rp.factory()
		if err != nil {
			rp.resources <- resourceWrapper{}
		}
	}
	return wrapper.resource, err
}

// Put will return a resource to the pool. For every successful Get,
// a corresponding Put is required. If you no longer need a resource,
// you will need to call Put(nil) instead of returning the closed resource.
// The will eventually cause a new resource to be created in its place.
func (rp *ResourcePool) Put(resource Resource) {
	var wrapper resourceWrapper
	if resource != nil {
		wrapper = resourceWrapper{resource, time.Now()}
	}
	select {
	case rp.resources <- wrapper:
	default:
		panic(errors.New("attempt to Put into a full ResourcePool"))
	}
}

// SetCapacity changes the capacity of the pool.
// You can use it to shrink or expand, but not beyond
// the max capacity. If the change requires the pool
// to be shrunk, SetCapacity waits till the necessary
// number of resources are returned to the pool.
// A SetCapacity of 0 is equivalent to closing the ResourcePool.
func (rp *ResourcePool) SetCapacity(capacity int) error {
	if capacity < 0 || capacity > cap(rp.resources) {
		return fmt.Errorf("capacity %d is out of range", capacity)
	}

	// Atomically swap new capacity with old, but only
	// if old capacity is non-zero.
	var oldcap int
	for {
		oldcap = int(rp.capacity.Get())
		if oldcap == 0 {
			return CLOSED_ERR
		}
		if oldcap == capacity {
			return nil
		}
		if rp.capacity.CompareAndSwap(int64(oldcap), int64(capacity)) {
			break
		}
	}

	if capacity < oldcap {
		for i := 0; i < oldcap-capacity; i++ {
			wrapper := <-rp.resources
			if wrapper.resource != nil {
				wrapper.resource.Close()
			}
		}
	} else {
		for i := 0; i < capacity-oldcap; i++ {
			rp.resources <- resourceWrapper{}
		}
	}
	if capacity == 0 {
		close(rp.resources)
	}
	return nil
}

func (rp *ResourcePool) recordWait(start time.Time) {
	rp.waitCount.Add(1)
	rp.waitTime.Add(time.Now().Sub(start))
}

func (rp *ResourcePool) SetIdleTimeout(idleTimeout time.Duration) {
	rp.idleTimeout.Set(idleTimeout)
}

func (rp *ResourcePool) StatsJSON() string {
	c, a, mx, wc, wt, it := rp.Stats()
	return fmt.Sprintf(`{"Capacity": %v, "Available": %v, "MaxCapacity": %v, "WaitCount": %v, "WaitTime": %v, "IdleTimeout": %v}`, c, a, mx, wc, int64(wt), int64(it))
}

func (rp *ResourcePool) Stats() (capacity, available, maxCap, waitCount int64, waitTime, idleTimeout time.Duration) {
	return rp.Capacity(), rp.Available(), rp.MaxCap(), rp.WaitCount(), rp.WaitTime(), rp.IdleTimeout()
}

func (rp *ResourcePool) Capacity() int64 {
	return rp.capacity.Get()
}

func (rp *ResourcePool) Available() int64 {
	return int64(len(rp.resources))
}

func (rp *ResourcePool) MaxCap() int64 {
	return int64(cap(rp.resources))
}

func (rp *ResourcePool) WaitCount() int64 {
	return rp.waitCount.Get()
}

func (rp *ResourcePool) WaitTime() time.Duration {
	return rp.waitTime.Get()
}

func (rp *ResourcePool) IdleTimeout() time.Duration {
	return rp.idleTimeout.Get()
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

package topo

import (
	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/trace"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/topo/topoproto"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// ShardReplicationInfo is the companion structure for ShardReplication.
type ShardReplicationInfo struct {
	*topodatapb.ShardReplication
	cell     string
	keyspace string
	shard    string
}

// NewShardReplicationInfo is for topo.Server implementations to
// create the structure
func NewShardReplicationInfo(sr *topodatapb.ShardReplication, cell, keyspace, shard string) *ShardReplicationInfo {
	return &ShardReplicationInfo{
		ShardReplication: sr,
		cell:             cell,
		keyspace:         keyspace,
		shard:            shard,
	}
}

// Cell returns the cell for a ShardReplicationInfo
func (sri *ShardReplicationInfo) Cell() string {
	return sri.cell
}

// Keyspace returns the keyspace for a ShardReplicationInfo
func (sri *ShardReplicationInfo) Keyspace() string {
	return sri.keyspace
}

// Shard returns the shard for a ShardReplicationInfo
func (sri *ShardReplicationInfo) Shard() string {
	return sri.shard
}

// GetShardReplicationNode finds a node for a given tablet.
func (sri *ShardReplicationInfo) GetShardReplicationNode(tabletAlias *topodatapb.TabletAlias) (*topodatapb.ShardReplication_Node, error) {
	for _, rl := range sri.Nodes {
		if *rl.TabletAlias == *tabletAlias {
			return rl, nil
		}
	}
	return nil, ErrNoNode
}

// UpdateShardReplicationRecord is a low level function to add / update an
// entry to the ShardReplication object.
func UpdateShardReplicationRecord(ctx context.Context, ts Server, keyspace, shard string, tabletAlias *topodatapb.TabletAlias) error {
	span := trace.NewSpanFromContext(ctx)
	span.StartClient("TopoServer.UpdateShardReplicationFields")
	span.Annotate("keyspace", keyspace)
	span.Annotate("shard", shard)
	span.Annotate("tablet", topoproto.TabletAliasString(tabletAlias))
	defer span.Finish()

	return ts.UpdateShardReplicationFields(ctx, tabletAlias.Cell, keyspace, shard, func(sr *topodatapb.ShardReplication) error {
		// Not very efficient, but easy to read, and allows us
		// to remove duplicate entries if any.
		nodes := make([]*topodatapb.ShardReplication_Node, 0, len(sr.Nodes)+1)
		found := false
		modified := false
		for _, node := range sr.Nodes {
			if *node.TabletAlias == *tabletAlias {
				if found {
					log.Warningf("Found a second ShardReplication_Node for tablet %v, deleting it", tabletAlias)
					modified = true
					continue
				}
				found = true
			}
			nodes = append(nodes, node)
		}
		if !found {
			nodes = append(nodes, &topodatapb.ShardReplication_Node{TabletAlias: tabletAlias})
			modified = true
		}
		if !modified {
			return ErrNoUpdateNeeded
		}
		sr.Nodes = nodes
		return nil
	})
}

// RemoveShardReplicationRecord is a low level function to remove an
// entry from the ShardReplication object.
func RemoveShardReplicationRecord(ctx context.Context, ts Server, cell, keyspace, shard string, tabletAlias *topodatapb.TabletAlias) error {
	err := ts.UpdateShardReplicationFields(ctx, cell, keyspace, shard, func(sr *topodatapb.ShardReplication) error {
		nodes := make([]*topodatapb.ShardReplication_Node, 0, len(sr.Nodes))
		for _, node := range sr.Nodes {
			if !proto.Equal(node.TabletAlias, tabletAlias) {
				nodes = append(nodes, node)
			}
		}
		sr.Nodes = nodes
		return nil
	})
	return err
}

// FixShardReplication will fix the first problem it encounters within
// a ShardReplication object.
func FixShardReplication(ctx context.Context, ts Server, logger logutil.Logger, cell, keyspace, shard string) error {
	sri, err := ts.GetShardReplication(ctx, cell, keyspace, shard)
	if err != nil {
		return err
	}

	for _, node := range sri.Nodes {
		ti, err := ts.GetTablet(ctx, node.TabletAlias)
		if err == ErrNoNode {
			logger.Warningf("Tablet %v is in the replication graph, but does not exist, removing it", node.TabletAlias)
			return RemoveShardReplicationRecord(ctx, ts, cell, keyspace, shard, node.TabletAlias)
		}
		if err != nil {
			// unknown error, we probably don't want to continue
			return err
		}

		if ti.Keyspace != keyspace || ti.Shard != shard || ti.Alias.Cell != cell {
			logger.Warningf("Tablet '%v' is in the replication graph, but has wrong keyspace/shard/cell, removing it", ti.Tablet)
			return RemoveShardReplicationRecord(ctx, ts, cell, keyspace, shard, node.TabletAlias)
		}

		logger.Infof("Keeping tablet %v in the replication graph", node.TabletAlias)
	}

	logger.Infof("All entries in replication graph are valid")
	return nil
}
<commit_msg>Fixing a couple proto compare.<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

package topo

import (
	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/trace"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/topo/topoproto"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// ShardReplicationInfo is the companion structure for ShardReplication.
type ShardReplicationInfo struct {
	*topodatapb.ShardReplication
	cell     string
	keyspace string
	shard    string
}

// NewShardReplicationInfo is for topo.Server implementations to
// create the structure
func NewShardReplicationInfo(sr *topodatapb.ShardReplication, cell, keyspace, shard string) *ShardReplicationInfo {
	return &ShardReplicationInfo{
		ShardReplication: sr,
		cell:             cell,
		keyspace:         keyspace,
		shard:            shard,
	}
}

// Cell returns the cell for a ShardReplicationInfo
func (sri *ShardReplicationInfo) Cell() string {
	return sri.cell
}

// Keyspace returns the keyspace for a ShardReplicationInfo
func (sri *ShardReplicationInfo) Keyspace() string {
	return sri.keyspace
}

// Shard returns the shard for a ShardReplicationInfo
func (sri *ShardReplicationInfo) Shard() string {
	return sri.shard
}

// GetShardReplicationNode finds a node for a given tablet.
func (sri *ShardReplicationInfo) GetShardReplicationNode(tabletAlias *topodatapb.TabletAlias) (*topodatapb.ShardReplication_Node, error) {
	for _, rl := range sri.Nodes {
		if proto.Equal(rl.TabletAlias, tabletAlias) {
			return rl, nil
		}
	}
	return nil, ErrNoNode
}

// UpdateShardReplicationRecord is a low level function to add / update an
// entry to the ShardReplication object.
func UpdateShardReplicationRecord(ctx context.Context, ts Server, keyspace, shard string, tabletAlias *topodatapb.TabletAlias) error {
	span := trace.NewSpanFromContext(ctx)
	span.StartClient("TopoServer.UpdateShardReplicationFields")
	span.Annotate("keyspace", keyspace)
	span.Annotate("shard", shard)
	span.Annotate("tablet", topoproto.TabletAliasString(tabletAlias))
	defer span.Finish()

	return ts.UpdateShardReplicationFields(ctx, tabletAlias.Cell, keyspace, shard, func(sr *topodatapb.ShardReplication) error {
		// Not very efficient, but easy to read, and allows us
		// to remove duplicate entries if any.
		nodes := make([]*topodatapb.ShardReplication_Node, 0, len(sr.Nodes)+1)
		found := false
		modified := false
		for _, node := range sr.Nodes {
			if proto.Equal(node.TabletAlias, tabletAlias) {
				if found {
					log.Warningf("Found a second ShardReplication_Node for tablet %v, deleting it", tabletAlias)
					modified = true
					continue
				}
				found = true
			}
			nodes = append(nodes, node)
		}
		if !found {
			nodes = append(nodes, &topodatapb.ShardReplication_Node{TabletAlias: tabletAlias})
			modified = true
		}
		if !modified {
			return ErrNoUpdateNeeded
		}
		sr.Nodes = nodes
		return nil
	})
}

// RemoveShardReplicationRecord is a low level function to remove an
// entry from the ShardReplication object.
func RemoveShardReplicationRecord(ctx context.Context, ts Server, cell, keyspace, shard string, tabletAlias *topodatapb.TabletAlias) error {
	err := ts.UpdateShardReplicationFields(ctx, cell, keyspace, shard, func(sr *topodatapb.ShardReplication) error {
		nodes := make([]*topodatapb.ShardReplication_Node, 0, len(sr.Nodes))
		for _, node := range sr.Nodes {
			if !proto.Equal(node.TabletAlias, tabletAlias) {
				nodes = append(nodes, node)
			}
		}
		sr.Nodes = nodes
		return nil
	})
	return err
}

// FixShardReplication will fix the first problem it encounters within
// a ShardReplication object.
func FixShardReplication(ctx context.Context, ts Server, logger logutil.Logger, cell, keyspace, shard string) error {
	sri, err := ts.GetShardReplication(ctx, cell, keyspace, shard)
	if err != nil {
		return err
	}

	for _, node := range sri.Nodes {
		ti, err := ts.GetTablet(ctx, node.TabletAlias)
		if err == ErrNoNode {
			logger.Warningf("Tablet %v is in the replication graph, but does not exist, removing it", node.TabletAlias)
			return RemoveShardReplicationRecord(ctx, ts, cell, keyspace, shard, node.TabletAlias)
		}
		if err != nil {
			// unknown error, we probably don't want to continue
			return err
		}

		if ti.Keyspace != keyspace || ti.Shard != shard || ti.Alias.Cell != cell {
			logger.Warningf("Tablet '%v' is in the replication graph, but has wrong keyspace/shard/cell, removing it", ti.Tablet)
			return RemoveShardReplicationRecord(ctx, ts, cell, keyspace, shard, node.TabletAlias)
		}

		logger.Infof("Keeping tablet %v in the replication graph", node.TabletAlias)
	}

	logger.Infof("All entries in replication graph are valid")
	return nil
}
<|endoftext|>


<commit_before>// Package vitessdriver contains the Vitess Go SQL driver.
/*
Vitess

Vitess is an SQL middleware which turns MySQL/MariaDB into a fast, scalable and
highly-available distributed database.
For more information, visit http://www.vitess.io.

Example

Using this SQL driver is as simple as:

  import (
    "time"
    "github.com/youtube/vitess/go/vt/vitessdriver"
  )

  func main() {
    // Connect to vtgate (v1 mode).
    db, err := vitessdriver.Open("localhost:15991", "keyspace", "0", "master", 30*time.Second)

    // Use "db" via the Golang sql interface.
  }

For a full example, please see: https://github.com/youtube/vitess/blob/master/examples/local/client.go

The full example is based on our tutorial for running Vitess locally: http://vitess.io/getting-started/local-instance.html

Vtgate

This driver connects to a vtgate process. In Vitess, vtgate is a proxy server
that routes queries to the appropriate shards.

By decoupling the routing logic from the application, vtgate enables Vitess
features like consistent, application-transparent resharding:
When you scale up the number of shards, vtgate becomes aware of the new shards.
You do not have to update your application at all.

Another Vitess feature is the vtgate v3 API. Once enabled, a sharded Vitess
setup becomes a drop-in replacement for an unsharded MySQL/MariaDB setup.

Vtgate API Versions

This driver supports the vtgate API versions v1 and v3, but not v2.

We recommend you use the v1 API for initial experiments and then switch to v3
later because v3 requires an additional setup step.

The API versions differ as follows:

Vtgate API v1

v1 requires each query to specify the target shard. vtgate then routes the query
accordingly.

This driver supports the v1 API at the connection level. For example, to open a
Vitess database connection for queries targeting shard "0" in the keyspace
"test_keyspace", call:

  vitessdriver.OpenShard(vtgateAddress, "test_keyspace", "0", 30*time.Second)

This mode is recommended for initial experiments or when migrating from a
pure MySQL/MariaDB setup to an unsharded Vitess setup. For a sharded Vitess
setup we recommend to move to the v2 or v3 API because only these APIs are
agnostic of resharding events.

Vtgate API v2

v2 requires each query to specify the affected sharding keys ("keyspace_id").

Based on the provided sharding key, vtgate then routes the query to the shard
that currently covers that sharding key.

This driver currently does not support the v2 API because it's difficult to
enhance the Go SQL interface with the sharding key information at the
query-level.
Instead, it's recommended to use the v3 API.

Vtgate API v3

v3 automatically infers the sharding key (and therefore the target shard)
from the query itself.

It also maintains secondary indexes across shards and generates values for
AUTO_INCREMENT columns.

See the vtgate v3 Features doc for an overview:
https://github.com/youtube/vitess/blob/master/doc/VTGateV3Features.md

To enable vtgate v3, you need to create a VSchema. A VSchema defines for vtgate
the properties of your Vitess setup (e.g. what is the sharding key, which
cross-shard indexes should be maintained, what are the AUTO_INCREMENT columns.)

As of 12/2015, the VSchema creation is not documented yet as we are in the
process of simplifying the VSchema definition and the overall process for
creating one.
If you want to create your own VSchema, we recommend to have a
look at the VSchema from the vtgate v3 demo:
https://github.com/youtube/vitess/blob/master/examples/demo/schema/vschema.json

(The demo itself is interactive and can be run by executing "./run.py" in the
"examples/demo/" directory.)

The vtgate v3 design doc, which we will also update and simplify in the future,
contains more details on the VSchema:
https://github.com/youtube/vitess/blob/master/doc/VTGateV3Design.md
*/
package vitessdriver
<commit_msg>vitessdriver: Addressed reviewer comments on godoc (second round).<commit_after>// Package vitessdriver contains the Vitess Go SQL driver.
/*
Vitess

Vitess is an SQL middleware which turns MySQL/MariaDB into a fast, scalable and
highly-available distributed database.
For more information, visit http://www.vitess.io.

Example

Using this SQL driver is as simple as:

  import (
    "time"
    "github.com/youtube/vitess/go/vt/vitessdriver"
  )

  func main() {
    // Connect to vtgate (v1 mode).
    db, err := vitessdriver.Open("localhost:15991", "keyspace", "0", "master", 30*time.Second)

    // Use "db" via the Golang sql interface.
  }

For a full example, please see: https://github.com/youtube/vitess/blob/master/examples/local/client.go

The full example is based on our tutorial for running Vitess locally: http://vitess.io/getting-started/local-instance.html

Vtgate

This driver connects to a vtgate process. In Vitess, vtgate is a proxy server
that routes queries to the appropriate shards.

By decoupling the routing logic from the application, vtgate enables Vitess
features like consistent, application-transparent resharding:
When you scale up the number of shards, vtgate becomes aware of the new shards.
You do not have to update your application at all.

Another Vitess feature is the vtgate v3 API. Once enabled, a sharded Vitess
setup becomes a drop-in replacement for an unsharded MySQL/MariaDB setup.

Vtgate API Versions

This driver supports the vtgate API versions v1 and v3, but not v2.

We recommend you use the v1 API for initial experiments and then switch to v3
later because v3 requires an additional setup step.

The API versions differ as follows:

Vtgate API v1

v1 requires each query to specify the target shard. vtgate then routes the query
accordingly.

This driver supports the v1 API at the connection level. For example, to open a
Vitess database connection for queries targeting shard "0" in the keyspace
"test_keyspace", call:

  vitessdriver.OpenShard(vtgateAddress, "test_keyspace", "0", 30*time.Second)

This mode is recommended for initial experiments or when migrating from a
pure MySQL/MariaDB setup to an unsharded Vitess setup. For a sharded Vitess
setup we recommend to move to the v2 or v3 API because only these APIs are
agnostic of resharding events.

Vtgate API v2

v2 requires each query to specify the affected sharding keys ("keyspace_id").

Based on the provided sharding key, vtgate then routes the query to the shard
that currently covers that sharding key.

This driver currently does not support the v2 API because it's difficult to
enhance the Go SQL interface with the sharding key information at the
query-level.
Instead, it's recommended to use the v3 API.

Vtgate API v3

v3 automatically infers the sharding key (and therefore the target shard)
from the query itself.

It also maintains secondary indexes across shards and generates values for
AUTO_INCREMENT columns.

See the vtgate v3 Features doc for an overview:
https://github.com/youtube/vitess/blob/master/doc/VTGateV3Features.md

To enable vtgate v3, you need to create a VSchema. A VSchema defines for vtgate
the properties of your Vitess setup (e.g. what is the sharding key, which
cross-shard indexes should be maintained, what are the AUTO_INCREMENT columns).

As of 12/2015, the VSchema creation is not documented yet as we are in the
process of simplifying the VSchema definition and the overall process for
creating one.
If you want to create your own VSchema, we recommend to have a
look at the VSchema from the vtgate v3 demo:
https://github.com/youtube/vitess/blob/master/examples/demo/schema/vschema.json

(The demo itself is interactive and can be run by executing "./run.py" in the
"examples/demo/" directory.)

The vtgate v3 design doc, which we will also update and simplify in the future,
contains more details on the VSchema:
https://github.com/youtube/vitess/blob/master/doc/VTGateV3Design.md
*/
package vitessdriver
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"os"
	"path/filepath"
	"github.com/kataras/iris/core/errors"
	"container/list"
	"sort"
)

func ls(p string) error {
	//Open dir and read
	if dir, err := os.Open(p); err == nil {
		if fileInfos, err := dir.Readdir(-1); err == nil {
			for _, fi := range fileInfos {
				fmt.Println(fi.Name())
			}
			dir.Close()
			return err
		}
		return errors.New("read dir failed")
	}
	return errors.New("open dir failed")
}

// sort demo
type IntArray []int

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

func (this IntArray) Less(i, j int) bool {
	return this[i] < this[j]
}

func (this IntArray) Swap(i, j int) {
	this[i], this[j] = this[j], this[i]
}

func main() {

	//sort
	arr2 := []int{3,54,7,843,2,87,54,1}
	sort.Ints(arr2)
	//sort.Sort(IntArray(arr2))
	fmt.Println(arr2)

	//Write file
	if rfile, err := os.Create("test.txt"); err == nil {
		rfile.WriteString("hello world")
		rfile.Close()
	}

	// Read file
	file, err := os.Open("test.txt")
	if err != nil {
		return
	}
	defer file.Close()

	//get file size
	if stat, err := file.Stat(); err == nil {
		bs := make([]byte, stat.Size())
		if _, err := file.Read(bs); err == nil {
			str := string(bs)
			fmt.Println(str)
		}
	}

	//walks the file tree and delete exe file
	filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		//fmt.Println(path)
		if !info.IsDir() && strings.HasSuffix(info.Name(), "exe") {
			os.Remove(path)
			fmt.Println(path)
			//return filepath.SkipDir
		}
		return nil
	})

	// ls
	err1 := ls(".")
	fmt.Println(err1)

	//strings
	//true 1 2 true false hello world, welcome world who-are-you [who are you] [97 98 99 49 50 51] hello
	//arr := []rune("test")
	arr := []byte("abc123")
	str := string([]byte{'h', 'e', 'l', 'l', 'o'})
	fmt.Println()
	fmt.Println(
		strings.Contains("test", "es"),
		strings.Index("test", "es"),
		strings.Count("test", "t"),
		strings.HasSuffix("package.go", "go"),
		strings.HasPrefix("package.go", "ack"),
		strings.Replace("hello pitou, welcome pitou", "pitou", "world", -1),
		strings.Join([]string{"who", "are", "you"}, "-"),
		strings.Split(strings.Join([]string{"who", "are", "you"}, "+"), "+"),
		arr, str,
	)

	//container
	var lst list.List
	lst.PushBack(1)
	lst.PushBack(1.12)
	lst.PushBack(112)
	lst.PushFront(100)
	for e := lst.Front(); e != nil; e = e.Next() {
		fmt.Println(e.Value)
	}
}
<commit_msg>feature: add packages demo<commit_after>package main

import (
	"fmt"
	"strings"
	"os"
	"path/filepath"
	"errors"
	"container/list"
	"sort"
)

func ls(p string) error {
	//Open dir and read
	if dir, err := os.Open(p); err == nil {
		if fileInfos, err := dir.Readdir(-1); err == nil {
			for _, fi := range fileInfos {
				fmt.Println(fi.Name())
			}
			dir.Close()
			return err
		}
		return errors.New("read dir failed")
	}
	return errors.New("open dir failed")
}

// sort demo
type IntArray []int

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

func (this IntArray) Less(i, j int) bool {
	return this[i] < this[j]
}

func (this IntArray) Swap(i, j int) {
	this[i], this[j] = this[j], this[i]
}

func main() {

	//sort
	arr2 := []int{3,54,7,843,2,87,54,1}
	sort.Ints(arr2)
	//sort.Sort(IntArray(arr2))
	fmt.Println(arr2)

	//Write file
	if rfile, err := os.Create("test.txt"); err == nil {
		rfile.WriteString("hello world")
		rfile.Close()
	}

	// Read file
	file, err := os.Open("test.txt")
	if err != nil {
		return
	}
	defer file.Close()

	//get file size
	if stat, err := file.Stat(); err == nil {
		bs := make([]byte, stat.Size())
		if _, err := file.Read(bs); err == nil {
			str := string(bs)
			fmt.Println(str)
		}
	}

	//walks the file tree and delete exe file
	filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		//fmt.Println(path)
		if !info.IsDir() && strings.HasSuffix(info.Name(), "exe") {
			os.Remove(path)
			fmt.Println(path)
			//return filepath.SkipDir
		}
		return nil
	})

	// ls
	err1 := ls(".")
	fmt.Println(err1)

	//strings
	//true 1 2 true false hello world, welcome world who-are-you [who are you] [97 98 99 49 50 51] hello
	//arr := []rune("test")
	arr := []byte("abc123")
	str := string([]byte{'h', 'e', 'l', 'l', 'o'})
	fmt.Println()
	fmt.Println(
		strings.Contains("test", "es"),
		strings.Index("test", "es"),
		strings.Count("test", "t"),
		strings.HasSuffix("package.go", "go"),
		strings.HasPrefix("package.go", "ack"),
		strings.Replace("hello pitou, welcome pitou", "pitou", "world", -1),
		strings.Join([]string{"who", "are", "you"}, "-"),
		strings.Split(strings.Join([]string{"who", "are", "you"}, "+"), "+"),
		arr, str,
	)

	//container
	var lst list.List
	lst.PushBack(1)
	lst.PushBack(1.12)
	lst.PushBack(112)
	lst.PushFront(100)
	for e := lst.Front(); e != nil; e = e.Next() {
		fmt.Println(e.Value)
	}
}
<|endoftext|>


<commit_before>package schema

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"sort"
)

var errInvalidIntervalzero = errors.New("interval cannot be 0")
var errInvalidOrgIdzero = errors.New("org-id cannot be 0")
var errInvalidEmptyName = errors.New("name cannot be empty")
var errInvalidMtype = errors.New("invalid mtype")
var errInvalidTagFormat = errors.New("invalid tag format")

type PartitionedMetric interface {
	Validate() error
	SetId()
	// return a []byte key comprised of the metric's OrgId
	// accepts an input []byte to allow callers to re-use
	// buffers to reduce memory allocations
	KeyByOrgId([]byte) []byte
	// return a []byte key comprised of the metric's Name
	// accepts an input []byte to allow callers to re-use
	// buffers to reduce memory allocations
	KeyBySeries([]byte) []byte
}

//go:generate msgp

// MetricData contains all metric metadata (some as fields, some as tags) and a datapoint
type MetricData struct {
	Id       string   `json:"id"`
	OrgId    int      `json:"org_id"`
	Name     string   `json:"name"`
	Interval int      `json:"interval"`
	Value    float64  `json:"value"`
	Unit     string   `json:"unit"`
	Time     int64    `json:"time"`
	Mtype    string   `json:"mtype"`
	Tags     []string `json:"tags" elastic:"type:string,index:not_analyzed"`
}

func (m *MetricData) Validate() error {
	if m.OrgId == 0 {
		return errInvalidOrgIdzero
	}
	if m.Interval == 0 {
		return errInvalidIntervalzero
	}
	if m.Name == "" {
		return errInvalidEmptyName
	}
	if m.Mtype == "" || (m.Mtype != "gauge" && m.Mtype != "rate" && m.Mtype != "count" && m.Mtype != "counter" && m.Mtype != "timestamp") {
		return errInvalidMtype
	}
	if !ValidateTags(m.Tags) {
		return errInvalidTagFormat
	}
	return nil
}

func (m *MetricData) KeyByOrgId(b []byte) []byte {
	if cap(b)-len(b) < 4 {
		// not enough unused space in the slice so we need to grow it.
		newBuf := make([]byte, len(b), len(b)+4)
		copy(newBuf, b)
		b = newBuf
	}
	// PutUint32 writes directly to the slice rather then appending.
	// so we need to set the length to 4 more bytes then it currently is.
	b = b[:len(b)+4]
	binary.LittleEndian.PutUint32(b[len(b)-4:], uint32(m.OrgId))
	return b
}

func (m *MetricData) KeyBySeries(b []byte) []byte {
	b = append(b, []byte(m.Name)...)
	return b
}

// returns a id (hash key) in the format OrgId.md5Sum
// the md5sum is a hash of the the concatination of the
// metric + each tag key:value pair (in metrics2.0 sense, so also fields), sorted alphabetically.
func (m *MetricData) SetId() {
	sort.Strings(m.Tags)

	buffer := bytes.NewBufferString(m.Name)
	buffer.WriteByte(0)
	buffer.WriteString(m.Unit)
	buffer.WriteByte(0)
	buffer.WriteString(m.Mtype)
	buffer.WriteByte(0)
	fmt.Fprintf(buffer, "%d", m.Interval)

	for _, k := range m.Tags {
		buffer.WriteByte(0)
		buffer.WriteString(k)
	}
	m.Id = fmt.Sprintf("%d.%x", m.OrgId, md5.Sum(buffer.Bytes()))
}

// can be used by some encoders, such as msgp
type MetricDataArray []*MetricData

// for ES
type MetricDefinition struct {
	Id       MKey   `json:"id"`
	OrgId    int    `json:"org_id"`
	Name     string `json:"name"`
	Interval int    `json:"interval"`                                      // minimum 10
	Unit     string `json:"unit"`
	Mtype    string `json:"mtype"`

	// some users of MetricDefinition (f.e. MetricTank) might add a "name" tag
	// to this slice which allows querying by name as a tag. this special tag
	// should not be stored or transmitted over the network, otherwise it may
	// just get overwritten by the receiver.
	Tags       []string `json:"tags" elastic:"type:string,index:not_analyzed"`
	LastUpdate int64    `json:"lastUpdate"` // unix timestamp
	Partition  int32    `json:"partition"`

	// this is a special attribute that does not need to be set, it is only used
	// to keep the state of NameWithTags()
	nameWithTags string `json:"-"`
}

// NameWithTags deduplicates the name and tags strings by storing their content
// as a single string in .nameWithTags and then makes .Name and the .Tags slices
// of it. Once this has been done it won't do it a second time, but just reuse
// the already generated .nameWithTags.
// It returns the deduplicated name with tags.
func (m *MetricDefinition) NameWithTags() string {
	if len(m.nameWithTags) > 0 {
		return m.nameWithTags
	}

	sort.Strings(m.Tags)

	nameWithTagsBuffer := bytes.NewBufferString(m.Name)
	tagPositions := make([]int, 0, len(m.Tags)*2)
	for _, t := range m.Tags {
		if len(t) >= 5 && t[:5] == "name=" {
			continue
		}

		nameWithTagsBuffer.WriteString(";")
		tagPositions = append(tagPositions, nameWithTagsBuffer.Len())
		nameWithTagsBuffer.WriteString(t)
		tagPositions = append(tagPositions, nameWithTagsBuffer.Len())
	}

	m.nameWithTags = nameWithTagsBuffer.String()
	m.Tags = make([]string, len(tagPositions)/2)
	for i := 0; i < len(m.Tags); i++ {
		m.Tags[i] = m.nameWithTags[tagPositions[i*2]:tagPositions[i*2+1]]
	}
	m.Name = m.nameWithTags[:len(m.Name)]

	return m.nameWithTags
}

func (m *MetricDefinition) SetId() {
	sort.Strings(m.Tags)

	buffer := bytes.NewBufferString(m.Name)
	buffer.WriteByte(0)
	buffer.WriteString(m.Unit)
	buffer.WriteByte(0)
	buffer.WriteString(m.Mtype)
	buffer.WriteByte(0)
	fmt.Fprintf(buffer, "%d", m.Interval)

	for _, t := range m.Tags {
		if len(t) >= 5 && t[:5] == "name=" {
			continue
		}

		buffer.WriteByte(0)
		buffer.WriteString(t)
	}

	m.Id = MKey{
		md5.Sum(buffer.Bytes()),
		uint32(m.OrgId),
	}
}

func (m *MetricDefinition) Validate() error {
	if m.OrgId == 0 {
		return errInvalidOrgIdzero
	}
	if m.Interval == 0 {
		return errInvalidIntervalzero
	}
	if m.Name == "" {
		return errInvalidEmptyName
	}
	if m.Mtype == "" || (m.Mtype != "gauge" && m.Mtype != "rate" && m.Mtype != "count" && m.Mtype != "counter" && m.Mtype != "timestamp") {
		return errInvalidMtype
	}
	if !ValidateTags(m.Tags) {
		return errInvalidTagFormat
	}
	return nil
}

func (m *MetricDefinition) KeyByOrgId(b []byte) []byte {
	if cap(b)-len(b) < 4 {
		// not enough unused space in the slice so we need to grow it.
		newBuf := make([]byte, len(b), len(b)+4)
		copy(newBuf, b)
		b = newBuf
	}
	// PutUint32 writes directly to the slice rather then appending.
	// so we need to set the length to 4 more bytes then it currently is.
	b = b[:len(b)+4]
	binary.LittleEndian.PutUint32(b[len(b)-4:], uint32(m.OrgId))
	return b
}

func (m *MetricDefinition) KeyBySeries(b []byte) []byte {
	b = append(b, []byte(m.Name)...)
	return b
}

// MetricDefinitionFromMetricData yields a MetricDefinition that has no references
// to the original MetricData
func MetricDefinitionFromMetricData(d *MetricData) *MetricDefinition {
	tags := make([]string, len(d.Tags))
	copy(tags, d.Tags)
	mkey, _ := MKeyFromString(d.Id)

	md := &MetricDefinition{
		Id:         mkey,
		Name:       d.Name,
		OrgId:      d.OrgId,
		Mtype:      d.Mtype,
		Interval:   d.Interval,
		LastUpdate: d.Time,
		Unit:       d.Unit,
		Tags:       tags,
	}

	return md
}

// ValidateTags returns whether all tags are in a valid format.
// a valid format is anything that looks like key=value,
// the length of key and value must be >0 and both must not contain the ; character.
func ValidateTags(tags []string) bool {
	for _, t := range tags {
		if len(t) == 0 {
			return false
		}

		if t[0] == '=' {
			return false
		}

		foundEqual := false
		for pos := 0; pos < len(t); pos++ {
			if t[pos] == ';' {
				return false
			}

			if !foundEqual {
				// no ! allowed in key
				if t[pos] == '!' {
					return false
				}

				// found the first =, so this will be the separator between key & value
				if t[pos] == '=' {
					// first equal sign must not be the last character
					if pos == len(t)-1 {
						return false
					}

					foundEqual = true
				}
			}
		}

		if !foundEqual {
			return false
		}
	}

	return true
}
<commit_msg>schema: clean up struct tags, remove references to elasticsearch<commit_after>package schema

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"sort"
)

var errInvalidIntervalzero = errors.New("interval cannot be 0")
var errInvalidOrgIdzero = errors.New("org-id cannot be 0")
var errInvalidEmptyName = errors.New("name cannot be empty")
var errInvalidMtype = errors.New("invalid mtype")
var errInvalidTagFormat = errors.New("invalid tag format")

type PartitionedMetric interface {
	Validate() error
	SetId()
	// return a []byte key comprised of the metric's OrgId
	// accepts an input []byte to allow callers to re-use
	// buffers to reduce memory allocations
	KeyByOrgId([]byte) []byte
	// return a []byte key comprised of the metric's Name
	// accepts an input []byte to allow callers to re-use
	// buffers to reduce memory allocations
	KeyBySeries([]byte) []byte
}

//go:generate msgp

// MetricData contains all metric metadata (some as fields, some as tags) and a datapoint
type MetricData struct {
	Id       string   `json:"id"`
	OrgId    int      `json:"org_id"`
	Name     string   `json:"name"`
	Interval int      `json:"interval"`
	Value    float64  `json:"value"`
	Unit     string   `json:"unit"`
	Time     int64    `json:"time"`
	Mtype    string   `json:"mtype"`
	Tags     []string `json:"tags"`
}

func (m *MetricData) Validate() error {
	if m.OrgId == 0 {
		return errInvalidOrgIdzero
	}
	if m.Interval == 0 {
		return errInvalidIntervalzero
	}
	if m.Name == "" {
		return errInvalidEmptyName
	}
	if m.Mtype == "" || (m.Mtype != "gauge" && m.Mtype != "rate" && m.Mtype != "count" && m.Mtype != "counter" && m.Mtype != "timestamp") {
		return errInvalidMtype
	}
	if !ValidateTags(m.Tags) {
		return errInvalidTagFormat
	}
	return nil
}

func (m *MetricData) KeyByOrgId(b []byte) []byte {
	if cap(b)-len(b) < 4 {
		// not enough unused space in the slice so we need to grow it.
		newBuf := make([]byte, len(b), len(b)+4)
		copy(newBuf, b)
		b = newBuf
	}
	// PutUint32 writes directly to the slice rather then appending.
	// so we need to set the length to 4 more bytes then it currently is.
	b = b[:len(b)+4]
	binary.LittleEndian.PutUint32(b[len(b)-4:], uint32(m.OrgId))
	return b
}

func (m *MetricData) KeyBySeries(b []byte) []byte {
	b = append(b, []byte(m.Name)...)
	return b
}

// returns a id (hash key) in the format OrgId.md5Sum
// the md5sum is a hash of the the concatination of the
// metric + each tag key:value pair (in metrics2.0 sense, so also fields), sorted alphabetically.
func (m *MetricData) SetId() {
	sort.Strings(m.Tags)

	buffer := bytes.NewBufferString(m.Name)
	buffer.WriteByte(0)
	buffer.WriteString(m.Unit)
	buffer.WriteByte(0)
	buffer.WriteString(m.Mtype)
	buffer.WriteByte(0)
	fmt.Fprintf(buffer, "%d", m.Interval)

	for _, k := range m.Tags {
		buffer.WriteByte(0)
		buffer.WriteString(k)
	}
	m.Id = fmt.Sprintf("%d.%x", m.OrgId, md5.Sum(buffer.Bytes()))
}

// can be used by some encoders, such as msgp
type MetricDataArray []*MetricData

type MetricDefinition struct {
	Id       MKey   `json:"mkey"`
	OrgId    int    `json:"org_id"`
	Name     string `json:"name"`
	Interval int    `json:"interval"`
	Unit     string `json:"unit"`
	Mtype    string `json:"mtype"`

	// some users of MetricDefinition (f.e. MetricTank) might add a "name" tag
	// to this slice which allows querying by name as a tag. this special tag
	// should not be stored or transmitted over the network, otherwise it may
	// just get overwritten by the receiver.
	Tags       []string `json:"tags"`
	LastUpdate int64    `json:"lastUpdate"` // unix timestamp
	Partition  int32    `json:"partition"`

	// this is a special attribute that does not need to be set, it is only used
	// to keep the state of NameWithTags()
	nameWithTags string `json:"-"`
}

// NameWithTags deduplicates the name and tags strings by storing their content
// as a single string in .nameWithTags and then makes .Name and the .Tags slices
// of it. Once this has been done it won't do it a second time, but just reuse
// the already generated .nameWithTags.
// It returns the deduplicated name with tags.
func (m *MetricDefinition) NameWithTags() string {
	if len(m.nameWithTags) > 0 {
		return m.nameWithTags
	}

	sort.Strings(m.Tags)

	nameWithTagsBuffer := bytes.NewBufferString(m.Name)
	tagPositions := make([]int, 0, len(m.Tags)*2)
	for _, t := range m.Tags {
		if len(t) >= 5 && t[:5] == "name=" {
			continue
		}

		nameWithTagsBuffer.WriteString(";")
		tagPositions = append(tagPositions, nameWithTagsBuffer.Len())
		nameWithTagsBuffer.WriteString(t)
		tagPositions = append(tagPositions, nameWithTagsBuffer.Len())
	}

	m.nameWithTags = nameWithTagsBuffer.String()
	m.Tags = make([]string, len(tagPositions)/2)
	for i := 0; i < len(m.Tags); i++ {
		m.Tags[i] = m.nameWithTags[tagPositions[i*2]:tagPositions[i*2+1]]
	}
	m.Name = m.nameWithTags[:len(m.Name)]

	return m.nameWithTags
}

func (m *MetricDefinition) SetId() {
	sort.Strings(m.Tags)

	buffer := bytes.NewBufferString(m.Name)
	buffer.WriteByte(0)
	buffer.WriteString(m.Unit)
	buffer.WriteByte(0)
	buffer.WriteString(m.Mtype)
	buffer.WriteByte(0)
	fmt.Fprintf(buffer, "%d", m.Interval)

	for _, t := range m.Tags {
		if len(t) >= 5 && t[:5] == "name=" {
			continue
		}

		buffer.WriteByte(0)
		buffer.WriteString(t)
	}

	m.Id = MKey{
		md5.Sum(buffer.Bytes()),
		uint32(m.OrgId),
	}
}

func (m *MetricDefinition) Validate() error {
	if m.OrgId == 0 {
		return errInvalidOrgIdzero
	}
	if m.Interval == 0 {
		return errInvalidIntervalzero
	}
	if m.Name == "" {
		return errInvalidEmptyName
	}
	if m.Mtype == "" || (m.Mtype != "gauge" && m.Mtype != "rate" && m.Mtype != "count" && m.Mtype != "counter" && m.Mtype != "timestamp") {
		return errInvalidMtype
	}
	if !ValidateTags(m.Tags) {
		return errInvalidTagFormat
	}
	return nil
}

func (m *MetricDefinition) KeyByOrgId(b []byte) []byte {
	if cap(b)-len(b) < 4 {
		// not enough unused space in the slice so we need to grow it.
		newBuf := make([]byte, len(b), len(b)+4)
		copy(newBuf, b)
		b = newBuf
	}
	// PutUint32 writes directly to the slice rather then appending.
	// so we need to set the length to 4 more bytes then it currently is.
	b = b[:len(b)+4]
	binary.LittleEndian.PutUint32(b[len(b)-4:], uint32(m.OrgId))
	return b
}

func (m *MetricDefinition) KeyBySeries(b []byte) []byte {
	b = append(b, []byte(m.Name)...)
	return b
}

// MetricDefinitionFromMetricData yields a MetricDefinition that has no references
// to the original MetricData
func MetricDefinitionFromMetricData(d *MetricData) *MetricDefinition {
	tags := make([]string, len(d.Tags))
	copy(tags, d.Tags)
	mkey, _ := MKeyFromString(d.Id)

	md := &MetricDefinition{
		Id:         mkey,
		Name:       d.Name,
		OrgId:      d.OrgId,
		Mtype:      d.Mtype,
		Interval:   d.Interval,
		LastUpdate: d.Time,
		Unit:       d.Unit,
		Tags:       tags,
	}

	return md
}

// ValidateTags returns whether all tags are in a valid format.
// a valid format is anything that looks like key=value,
// the length of key and value must be >0 and both must not contain the ; character.
func ValidateTags(tags []string) bool {
	for _, t := range tags {
		if len(t) == 0 {
			return false
		}

		if t[0] == '=' {
			return false
		}

		foundEqual := false
		for pos := 0; pos < len(t); pos++ {
			if t[pos] == ';' {
				return false
			}

			if !foundEqual {
				// no ! allowed in key
				if t[pos] == '!' {
					return false
				}

				// found the first =, so this will be the separator between key & value
				if t[pos] == '=' {
					// first equal sign must not be the last character
					if pos == len(t)-1 {
						return false
					}

					foundEqual = true
				}
			}
		}

		if !foundEqual {
			return false
		}
	}

	return true
}
<|endoftext|>


<commit_before>package hostdb

// scan.go contains the functions which periodically scan the list of all hosts
// to see which hosts are online or offline, and to get any updates to the
// settings of the hosts.

import (
	"bytes"
	"crypto/rand"
	"math/big"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	defaultScanSleep = 1*time.Hour + 37*time.Minute
	maxScanSleep     = 4 * time.Hour
	minScanSleep     = 1*time.Hour + 20*time.Minute

	maxActiveHosts              = 500
	inactiveHostCheckupQuantity = 1000

	maxSettingsLen = 2e3

	hostRequestTimeout = 30 * time.Second

	// scanningThreads is the number of threads that will be probing hosts for
	// their settings and checking for reliability.
	scanningThreads = 12
)

// Reliability is a measure of a host's uptime.
var (
	MaxReliability     = types.NewCurrency64(500) // Given the scanning defaults, about 6 weeks of survival.
	DefaultReliability = types.NewCurrency64(150) // Given the scanning defaults, about 2 week of survival.
	UnreachablePenalty = types.NewCurrency64(1)
)

// addHostToScanPool creates a gofunc that adds a host to the scan pool. If the
// scan pool is currently full, the blocking gofunc will not cause a deadlock.
// The gofunc is created inside of this function to eliminate the burden of
// needing to remember to call 'go addHostToScanPool'.
func (hdb *HostDB) scanHostEntry(entry *hostEntry) {
	// Add the entry to a waitlist, then check if any thread is currently
	// emptying the waitlist. If not, spawn a thread to empty the waitlist.
	hdb.scanList = append(hdb.scanList, entry)
	if hdb.scanWait {
		// Another thread is emptying the scan list, nothing to worry about.
		return
	}

	// Nobody is emptying the scan list, volunteer.
	hdb.scanWait = true
	go func() {
		if hdb.tg.Add() != nil {
			return
		}
		defer hdb.tg.Done()

		for {
			hdb.mu.Lock()
			if len(hdb.scanList) == 0 {
				// Scan list is empty, can exit. Let the world know that nobody
				// is emptying the scan list anymore.
				hdb.scanWait = false
				hdb.mu.Unlock()
				return
			}
			// Get the next host, shrink the scan list.
			entry := hdb.scanList[0]
			hdb.scanList = hdb.scanList[1:]
			hdb.mu.Unlock()

			// Block while we wait for an opening in the scan pool.
			select {
			case hdb.scanPool <- entry:
				// iterate again
			case <-hdb.tg.StopChan():
				// quit
				return
			}
		}
	}()
}

// decrementReliability reduces the reliability of a node, moving it out of the
// set of active hosts or deleting it entirely if necessary.
func (hdb *HostDB) decrementReliability(addr modules.NetAddress, penalty types.Currency) {
	hdb.log.Debugln("reliability decrement issued for", addr)

	// Look up the entry and decrement the reliability.
	entry, exists := hdb.allHosts[addr]
	if !exists {
		// TODO: should panic here
		return
	}
	entry.Reliability = entry.Reliability.Sub(penalty)
	entry.Online = false

	// If the entry is in the active database, remove it from the active
	// database.
	node, exists := hdb.activeHosts[addr]
	if exists {
		hdb.log.Debugln("host is being pulled from list of active hosts", addr)
		node.removeNode()
		delete(hdb.activeHosts, entry.NetAddress)
	}

	// If the reliability has fallen to 0, remove the host from the
	// database entirely.
	if entry.Reliability.IsZero() {
		hdb.log.Debugln("host is being dropped from hostdb", addr)
		delete(hdb.allHosts, addr)
	}
}

// managedUpdateEntry updates an entry in the hostdb after a scan has taken
// place.
func (hdb *HostDB) managedUpdateEntry(entry *hostEntry, newSettings modules.HostExternalSettings, netErr error) {
	hdb.mu.Lock()
	defer hdb.mu.Unlock()

	// Regardless of whether the host responded, add it to allHosts.
	priorHost, exists := hdb.allHosts[entry.NetAddress]
	if !exists {
		hdb.allHosts[entry.NetAddress] = entry
	}

	// If the scan was unsuccessful, decrement the host's reliability.
	if netErr != nil {
		if exists && bytes.Equal(priorHost.PublicKey.Key, entry.PublicKey.Key) {
			// Only decrement the reliability if the public key in the
			// hostdb matches the public key in the host announcement -
			// the failure may just be a failed signature, indicating
			// the wrong public key.
			hdb.decrementReliability(entry.NetAddress, UnreachablePenalty)
		}
		return
	}

	// The host entry should be updated to reflect the new weight. The safety
	// properties of the tree require that the weight does not change while the
	// node is in the tree, so the node must be removed before the settings and
	// weight are changed.
	existingNode, exists := hdb.activeHosts[entry.NetAddress]
	if exists {
		existingNode.removeNode()
		delete(hdb.activeHosts, entry.NetAddress)
	} else if len(hdb.activeHosts) > maxActiveHosts {
		// We already have the maximum number of active hosts, do not add more.
		return
	}

	// Update the host settings, reliability, and weight. The old NetAddress
	// must be preserved.
	newSettings.NetAddress = entry.HostExternalSettings.NetAddress
	entry.HostExternalSettings = newSettings
	entry.Reliability = MaxReliability
	entry.Online = true
	entry.Weight = calculateHostWeight(*entry)
	hdb.insertNode(entry)

	// Sanity check - the node should be in the hostdb now.
	_, exists = hdb.activeHosts[entry.NetAddress]
	if !exists {
		hdb.log.Critical("Host was not added to the list of active hosts after the entry was updated.")
	}
	hdb.save()
}

// threadedProbeHosts tries to fetch the settings of a host. If successful, the
// host is put in the set of active hosts. If unsuccessful, the host id deleted
// from the set of active hosts.
func (hdb *HostDB) threadedProbeHosts() {
	defer hdb.threadGroup.Done()
	for hostEntry := range hdb.scanPool {
		// Request settings from the queued host entry.
		// TODO: use dialer.Cancel to shutdown quickly
		//
		// A readlock is necessary when viewing the elements of the host entry.
		hdb.mu.RLock()
		netAddr := hostEntry.NetAddress
		pubKey := hostEntry.PublicKey
		hdb.mu.RUnlock()
		hdb.log.Debugln("Scanning", netAddr, pubKey)
		var settings modules.HostExternalSettings
		err := func() error {
			conn, err := hdb.dialer.DialTimeout(netAddr, hostRequestTimeout)
			if err != nil {
				return err
			}
			defer conn.Close()
			err = encoding.WriteObject(conn, modules.RPCSettings)
			if err != nil {
				return err
			}
			var pubkey crypto.PublicKey
			copy(pubkey[:], pubKey.Key)
			return crypto.ReadSignedObject(conn, &settings, maxSettingsLen, pubkey)
		}()
		if err != nil {
			hdb.log.Debugln("Scanning", netAddr, pubKey, "failed:", err)
		} else {
			hdb.log.Debugln("Scanning", netAddr, pubKey, "succeeded")
		}

		// Update the host tree to have a new entry.
		hdb.managedUpdateEntry(hostEntry, settings, err)
	}
}

// threadedScan is an ongoing function which will query the full set of hosts
// every few hours to see who is online and available for uploading.
func (hdb *HostDB) threadedScan() {
	defer hdb.threadGroup.Done()
	for {
		// Determine who to scan. At most 'maxActiveHosts' will be scanned,
		// starting with the active hosts followed by a random selection of the
		// inactive hosts.
		func() {
			hdb.mu.Lock()
			defer hdb.mu.Unlock()

			// Scan all active hosts.
			for _, host := range hdb.activeHosts {
				hdb.scanHostEntry(host.hostEntry)
			}

			// Assemble all of the inactive hosts into a single array.
			var entries []*hostEntry
			for _, entry := range hdb.allHosts {
				_, exists := hdb.activeHosts[entry.NetAddress]
				if !exists {
					entries = append(entries, entry)
				}
			}

			// Generate a random ordering of up to inactiveHostCheckupQuantity
			// hosts.
			hostOrder, err := crypto.Perm(len(entries))
			if err != nil {
				hdb.log.Println("ERR: could not generate random permutation:", err)
			}

			// Scan each host.
			for i := 0; i < len(hostOrder) && i < inactiveHostCheckupQuantity; i++ {
				hdb.scanHostEntry(entries[hostOrder[i]])
			}
		}()

		// Sleep for a random amount of time before doing another round of
		// scanning. The minimums and maximums keep the scan time reasonable,
		// while the randomness prevents the scanning from always happening at
		// the same time of day or week.
		maxBig := big.NewInt(int64(maxScanSleep))
		minBig := big.NewInt(int64(minScanSleep))
		randSleep, err := rand.Int(rand.Reader, maxBig.Sub(maxBig, minBig))
		if err != nil {
			build.Critical(err)
			// If there's an error, sleep for the default amount of time.
			defaultBig := big.NewInt(int64(defaultScanSleep))
			randSleep = defaultBig.Sub(defaultBig, minBig)
		}

		select {
		// awaken and exit if hostdb is closing
		case <-hdb.closeChan:
			return
		case <-time.After(time.Duration(randSleep.Int64()) + minScanSleep):
		}
	}
}
<commit_msg>resolve hostdb thread group edge case<commit_after>package hostdb

// scan.go contains the functions which periodically scan the list of all hosts
// to see which hosts are online or offline, and to get any updates to the
// settings of the hosts.

import (
	"bytes"
	"crypto/rand"
	"math/big"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	defaultScanSleep = 1*time.Hour + 37*time.Minute
	maxScanSleep     = 4 * time.Hour
	minScanSleep     = 1*time.Hour + 20*time.Minute

	maxActiveHosts              = 500
	inactiveHostCheckupQuantity = 1000

	maxSettingsLen = 2e3

	hostRequestTimeout = 30 * time.Second

	// scanningThreads is the number of threads that will be probing hosts for
	// their settings and checking for reliability.
	scanningThreads = 12
)

// Reliability is a measure of a host's uptime.
var (
	MaxReliability     = types.NewCurrency64(500) // Given the scanning defaults, about 6 weeks of survival.
	DefaultReliability = types.NewCurrency64(150) // Given the scanning defaults, about 2 week of survival.
	UnreachablePenalty = types.NewCurrency64(1)
)

// addHostToScanPool creates a gofunc that adds a host to the scan pool. If the
// scan pool is currently full, the blocking gofunc will not cause a deadlock.
// The gofunc is created inside of this function to eliminate the burden of
// needing to remember to call 'go addHostToScanPool'.
func (hdb *HostDB) scanHostEntry(entry *hostEntry) {
	// Add the entry to a waitlist, then check if any thread is currently
	// emptying the waitlist. If not, spawn a thread to empty the waitlist.
	hdb.scanList = append(hdb.scanList, entry)
	if hdb.scanWait {
		// Another thread is emptying the scan list, nothing to worry about.
		return
	}

	// Nobody is emptying the scan list, volunteer.
	if hdb.tg.Add() != nil {
		// Hostdb is shutting down, don't spin up another thread.
		return
	}
	hdb.scanWait = true
	go func() {
		defer hdb.tg.Done()

		for {
			hdb.mu.Lock()
			if len(hdb.scanList) == 0 {
				// Scan list is empty, can exit. Let the world know that nobody
				// is emptying the scan list anymore.
				hdb.scanWait = false
				hdb.mu.Unlock()
				return
			}
			// Get the next host, shrink the scan list.
			entry := hdb.scanList[0]
			hdb.scanList = hdb.scanList[1:]
			hdb.mu.Unlock()

			// Block while we wait for an opening in the scan pool.
			select {
			case hdb.scanPool <- entry:
				// iterate again
			case <-hdb.tg.StopChan():
				// quit
				return
			}
		}
	}()
}

// decrementReliability reduces the reliability of a node, moving it out of the
// set of active hosts or deleting it entirely if necessary.
func (hdb *HostDB) decrementReliability(addr modules.NetAddress, penalty types.Currency) {
	hdb.log.Debugln("reliability decrement issued for", addr)

	// Look up the entry and decrement the reliability.
	entry, exists := hdb.allHosts[addr]
	if !exists {
		// TODO: should panic here
		return
	}
	entry.Reliability = entry.Reliability.Sub(penalty)
	entry.Online = false

	// If the entry is in the active database, remove it from the active
	// database.
	node, exists := hdb.activeHosts[addr]
	if exists {
		hdb.log.Debugln("host is being pulled from list of active hosts", addr)
		node.removeNode()
		delete(hdb.activeHosts, entry.NetAddress)
	}

	// If the reliability has fallen to 0, remove the host from the
	// database entirely.
	if entry.Reliability.IsZero() {
		hdb.log.Debugln("host is being dropped from hostdb", addr)
		delete(hdb.allHosts, addr)
	}
}

// managedUpdateEntry updates an entry in the hostdb after a scan has taken
// place.
func (hdb *HostDB) managedUpdateEntry(entry *hostEntry, newSettings modules.HostExternalSettings, netErr error) {
	hdb.mu.Lock()
	defer hdb.mu.Unlock()

	// Regardless of whether the host responded, add it to allHosts.
	priorHost, exists := hdb.allHosts[entry.NetAddress]
	if !exists {
		hdb.allHosts[entry.NetAddress] = entry
	}

	// If the scan was unsuccessful, decrement the host's reliability.
	if netErr != nil {
		if exists && bytes.Equal(priorHost.PublicKey.Key, entry.PublicKey.Key) {
			// Only decrement the reliability if the public key in the
			// hostdb matches the public key in the host announcement -
			// the failure may just be a failed signature, indicating
			// the wrong public key.
			hdb.decrementReliability(entry.NetAddress, UnreachablePenalty)
		}
		return
	}

	// The host entry should be updated to reflect the new weight. The safety
	// properties of the tree require that the weight does not change while the
	// node is in the tree, so the node must be removed before the settings and
	// weight are changed.
	existingNode, exists := hdb.activeHosts[entry.NetAddress]
	if exists {
		existingNode.removeNode()
		delete(hdb.activeHosts, entry.NetAddress)
	} else if len(hdb.activeHosts) > maxActiveHosts {
		// We already have the maximum number of active hosts, do not add more.
		return
	}

	// Update the host settings, reliability, and weight. The old NetAddress
	// must be preserved.
	newSettings.NetAddress = entry.HostExternalSettings.NetAddress
	entry.HostExternalSettings = newSettings
	entry.Reliability = MaxReliability
	entry.Online = true
	entry.Weight = calculateHostWeight(*entry)
	hdb.insertNode(entry)

	// Sanity check - the node should be in the hostdb now.
	_, exists = hdb.activeHosts[entry.NetAddress]
	if !exists {
		hdb.log.Critical("Host was not added to the list of active hosts after the entry was updated.")
	}
	hdb.save()
}

// threadedProbeHosts tries to fetch the settings of a host. If successful, the
// host is put in the set of active hosts. If unsuccessful, the host id deleted
// from the set of active hosts.
func (hdb *HostDB) threadedProbeHosts() {
	defer hdb.threadGroup.Done()
	for hostEntry := range hdb.scanPool {
		// Request settings from the queued host entry.
		// TODO: use dialer.Cancel to shutdown quickly
		//
		// A readlock is necessary when viewing the elements of the host entry.
		hdb.mu.RLock()
		netAddr := hostEntry.NetAddress
		pubKey := hostEntry.PublicKey
		hdb.mu.RUnlock()
		hdb.log.Debugln("Scanning", netAddr, pubKey)
		var settings modules.HostExternalSettings
		err := func() error {
			conn, err := hdb.dialer.DialTimeout(netAddr, hostRequestTimeout)
			if err != nil {
				return err
			}
			defer conn.Close()
			err = encoding.WriteObject(conn, modules.RPCSettings)
			if err != nil {
				return err
			}
			var pubkey crypto.PublicKey
			copy(pubkey[:], pubKey.Key)
			return crypto.ReadSignedObject(conn, &settings, maxSettingsLen, pubkey)
		}()
		if err != nil {
			hdb.log.Debugln("Scanning", netAddr, pubKey, "failed:", err)
		} else {
			hdb.log.Debugln("Scanning", netAddr, pubKey, "succeeded")
		}

		// Update the host tree to have a new entry.
		hdb.managedUpdateEntry(hostEntry, settings, err)
	}
}

// threadedScan is an ongoing function which will query the full set of hosts
// every few hours to see who is online and available for uploading.
func (hdb *HostDB) threadedScan() {
	defer hdb.threadGroup.Done()
	for {
		// Determine who to scan. At most 'maxActiveHosts' will be scanned,
		// starting with the active hosts followed by a random selection of the
		// inactive hosts.
		func() {
			hdb.mu.Lock()
			defer hdb.mu.Unlock()

			// Scan all active hosts.
			for _, host := range hdb.activeHosts {
				hdb.scanHostEntry(host.hostEntry)
			}

			// Assemble all of the inactive hosts into a single array.
			var entries []*hostEntry
			for _, entry := range hdb.allHosts {
				_, exists := hdb.activeHosts[entry.NetAddress]
				if !exists {
					entries = append(entries, entry)
				}
			}

			// Generate a random ordering of up to inactiveHostCheckupQuantity
			// hosts.
			hostOrder, err := crypto.Perm(len(entries))
			if err != nil {
				hdb.log.Println("ERR: could not generate random permutation:", err)
			}

			// Scan each host.
			for i := 0; i < len(hostOrder) && i < inactiveHostCheckupQuantity; i++ {
				hdb.scanHostEntry(entries[hostOrder[i]])
			}
		}()

		// Sleep for a random amount of time before doing another round of
		// scanning. The minimums and maximums keep the scan time reasonable,
		// while the randomness prevents the scanning from always happening at
		// the same time of day or week.
		maxBig := big.NewInt(int64(maxScanSleep))
		minBig := big.NewInt(int64(minScanSleep))
		randSleep, err := rand.Int(rand.Reader, maxBig.Sub(maxBig, minBig))
		if err != nil {
			build.Critical(err)
			// If there's an error, sleep for the default amount of time.
			defaultBig := big.NewInt(int64(defaultScanSleep))
			randSleep = defaultBig.Sub(defaultBig, minBig)
		}

		select {
		// awaken and exit if hostdb is closing
		case <-hdb.closeChan:
			return
		case <-time.After(time.Duration(randSleep.Int64()) + minScanSleep):
		}
	}
}
<|endoftext|>