<commit_before>// Things not yet implemented:
// - enforcing use="restricted" on attributes
// - namespaces

package main

import (
	"fmt"
	"log"
	"os"
	"strings"
)

var (
	usage = `Usage:

  goxsd <xsd>

Arguments:

  xsd     Path to a valid XSD file

goxsd is a tool for generating XML decoding Go structs, according to an XSD
schema.

The argument is expected to be the path to a valid XSD schema file. Any import
statements in that file will be be followed and parsed. The resulting set of
Go structs will be printed on stdout.
`
)

func main() {
	if len(os.Args) != 2 {
		fmt.Println(usage)
		os.Exit(1)
	}
	xsdFile := os.Args[1]

	s, err := extractSchemas(xsdFile)
	if err != nil {
		log.Fatal(err)
	}
	builder := newBuilder(s)
	parse(os.Stdout, builder.buildXML())
}

type xmlElem struct {
	Name     string
	Type     string
	List     bool
	Cdata    bool
	Attribs  []xmlAttrib
	Children []*xmlElem
}

type xmlAttrib struct {
	Name string
	Type string
}

type builder struct {
	schemas    []xsdSchema
	complTypes map[string]xsdComplexType
	simplTypes map[string]xsdSimpleType
}

func newBuilder(s []xsdSchema) builder {
	return builder{
		schemas:    s,
		complTypes: make(map[string]xsdComplexType),
		simplTypes: make(map[string]xsdSimpleType),
	}
}

func (b builder) buildXML() []*xmlElem {
	var roots []xsdElement
	for _, s := range b.schemas {
		for _, e := range s.Elements {
			roots = append(roots, e)
		}
		for _, t := range s.ComplexTypes {
			b.complTypes[t.Name] = t
		}
		for _, t := range s.SimpleTypes {
			b.simplTypes[t.Name] = t
		}
	}

	var xelems []*xmlElem
	for _, e := range roots {
		xelems = append(xelems, b.buildFromElement(e))
	}

	return xelems
}

// Build a xmlElem from an xsdElement, recursively traverse the XSD type
// information to build up a XML descendant hierarchy.
func (b builder) buildFromElement(e xsdElement) *xmlElem {
	xelem := &xmlElem{Name: e.Name, Type: e.Name}

	if e.isList() {
		xelem.List = true
	}

	if !e.inlineType() {
		switch t := b.findType(e.Type).(type) {
		case xsdComplexType:
			b.buildFromComplexType(xelem, t)
		case xsdSimpleType:
			buildFromSimpleType(xelem, t)
		case string:
			xelem.Type = t
		}
		return xelem
	}

	if e.ComplexType != nil { // inline complex type
		b.buildFromComplexType(xelem, *e.ComplexType)
		return xelem
	}

	if e.SimpleType != nil { // inline simple type
		buildFromSimpleType(xelem, *e.SimpleType)
		return xelem
	}

	return xelem
}

func (b builder) buildFromComplexContent(xelem *xmlElem, c xsdComplexContent) {
	if c.Extension != nil {
		if c.Extension.Sequence != nil {
			for _, e := range c.Extension.Sequence {
				xelem.Children = append(xelem.Children, b.buildFromElement(e))
			}
		}
		base := c.Extension.Base
		switch t := b.findType(base).(type) {
		case xsdComplexType:
			b.buildFromComplexType(xelem, t)
		}

	}
}

func typeFromXsdType(typ string) string {
	switch typ {
	case "boolean":
		typ = "bool"
	case "language", "dateTime", "Name", "token":
		typ = "string"
	case "long", "short", "integer", "int":
		typ = "int"
	case "decimal":
		typ = "float64"
	}
	return typ
}

func addAttributes(xelem *xmlElem, attribs []xsdAttribute) {
	if attribs != nil {
		for _, attr := range attribs {
			typ := typeFromXsdType(stripNamespace(attr.Type))
			xelem.Attribs = append(xelem.Attribs, xmlAttrib{Name: attr.Name, Type: typ})
		}
	}
}

// A simple content can refer to a text-only complex type
func (b builder) buildFromSimpleContent(xelem *xmlElem, c xsdSimpleContent) {
	if c.Extension != nil {
		// (annotation?, ((group|all|choice|sequence)?, ((attribute|attributeGroup)*, anyAttribute?)))
		if c.Extension.Attributes != nil {
			b.buildFromAttributes(xelem, c.Extension.Attributes)
		}
		// has always a base type

		var child *xmlElem
		switch t := b.findType(c.Extension.Base).(type) {
		case xsdComplexType:
			b.buildFromComplexType(xelem, t)
		case xsdSimpleType:
			child = &xmlElem{Name: xelem.Name, Cdata: true}
			buildFromSimpleType(child, t)
			xelem.Children = []*xmlElem{child}
		default:
			child = &xmlElem{Name: xelem.Name, Cdata: true}
			child.Type = typeFromXsdType(t.(string))
			xelem.Children = []*xmlElem{child}
		}
	}

	if c.Restriction != nil {
		switch t := b.findType(c.Restriction.Base).(type) {
		case xsdComplexType:
			b.buildFromComplexType(xelem, t)
		case xsdSimpleType:
			buildFromSimpleType(xelem, t)
		default:
			xelem.Type = typeFromXsdType(t.(string))
			//addAttributes(xelem, c.Extension.Attributes)
		}
	}
}

func (b builder) buildFromAttributes(xelem *xmlElem, attrs []xsdAttribute) {
	for _, a := range attrs {
		attr := xmlAttrib{Name: a.Name}
		switch t := b.findType(a.Type).(type) {
		case xsdSimpleType:
			// Get type name from simpleType
			// If Restriction.Base is a simpleType or complexType, we panic
			attr.Type = b.findType(t.Restriction.Base).(string)
		case string:
			// If empty, then simpleType is present as content, but we ignore
			// that now
			attr.Type = t
		}
		xelem.Attribs = append(xelem.Attribs, attr)
	}
}

// buildFromComplexType takes an xmlElem and an xsdComplexType, containing
// XSD type information for xmlElem enrichment.
func (b builder) buildFromComplexType(xelem *xmlElem, t xsdComplexType) {
	if t.Sequence != nil { // Does the element have children?
		for _, e := range t.Sequence {
			xelem.Children = append(xelem.Children, b.buildFromElement(e))
		}
	}

	if t.Attributes != nil {
		b.buildFromAttributes(xelem, t.Attributes)
	}

	if t.ComplexContent != nil {
		b.buildFromComplexContent(xelem, *t.ComplexContent)
	}

	if t.SimpleContent != nil {
		b.buildFromSimpleContent(xelem, *t.SimpleContent)
	}
}

func buildFromSimpleType(xelem *xmlElem, t xsdSimpleType) {
	xelem.Type = typeFromXsdType(stripNamespace(t.Restriction.Base))
}

// findType takes a type name and checks if it is a registered XSD type
// (simple or complex), in which case that type is returned. If no such
// type can be found, the XSD specific primitive types are mapped to their
// Go correspondents. If no XSD type was found, the type name itself is
// returned.
func (b builder) findType(name string) interface{} {
	name = stripNamespace(name)
	if t, ok := b.complTypes[name]; ok {
		return t
	}
	if t, ok := b.simplTypes[name]; ok {
		return t
	}

	switch name {
	case "boolean":
		return "bool"
	case "language", "dateTime", "Name", "token":
		return "string"
	case "long", "short", "integer", "int":
		return "int"
	case "decimal":
		return "float64"
	default:
		return name
	}
}

func stripNamespace(name string) string {
	if s := strings.Split(name, ":"); len(s) > 1 {
		return s[len(s)-1]
	}
	return name
}
<commit_msg>Refactor<commit_after>// Things not yet implemented:
// - enforcing use="restricted" on attributes
// - namespaces

package main

import (
	"fmt"
	"log"
	"os"
	"strings"
)

var (
	usage = `Usage:

  goxsd <xsd>

Arguments:

  xsd     Path to a valid XSD file

goxsd is a tool for generating XML decoding Go structs, according to an XSD
schema.

The argument is expected to be the path to a valid XSD schema file. Any import
statements in that file will be be followed and parsed. The resulting set of
Go structs will be printed on stdout.
`
)

func main() {
	if len(os.Args) != 2 {
		fmt.Println(usage)
		os.Exit(1)
	}
	xsdFile := os.Args[1]

	s, err := extractSchemas(xsdFile)
	if err != nil {
		log.Fatal(err)
	}
	builder := newBuilder(s)
	parse(os.Stdout, builder.buildXML())
}

type xmlElem struct {
	Name     string
	Type     string
	List     bool
	Cdata    bool
	Attribs  []xmlAttrib
	Children []*xmlElem
}

type xmlAttrib struct {
	Name string
	Type string
}

type builder struct {
	schemas    []xsdSchema
	complTypes map[string]xsdComplexType
	simplTypes map[string]xsdSimpleType
}

func newBuilder(s []xsdSchema) builder {
	return builder{
		schemas:    s,
		complTypes: make(map[string]xsdComplexType),
		simplTypes: make(map[string]xsdSimpleType),
	}
}

func (b builder) buildXML() []*xmlElem {
	var roots []xsdElement
	for _, s := range b.schemas {
		for _, e := range s.Elements {
			roots = append(roots, e)
		}
		for _, t := range s.ComplexTypes {
			b.complTypes[t.Name] = t
		}
		for _, t := range s.SimpleTypes {
			b.simplTypes[t.Name] = t
		}
	}

	var xelems []*xmlElem
	for _, e := range roots {
		xelems = append(xelems, b.buildFromElement(e))
	}

	return xelems
}

// buildFromElement builds an xmlElem from an xsdElement, recursively
// traversing the XSD type information to build up an XML element hierarchy.
func (b builder) buildFromElement(e xsdElement) *xmlElem {
	xelem := &xmlElem{Name: e.Name, Type: e.Name}

	if e.isList() {
		xelem.List = true
	}

	if !e.inlineType() {
		switch t := b.findType(e.Type).(type) {
		case xsdComplexType:
			b.buildFromComplexType(xelem, t)
		case xsdSimpleType:
			b.buildFromSimpleType(xelem, t)
		case string:
			xelem.Type = t
		}
		return xelem
	}

	if e.ComplexType != nil { // inline complex type
		b.buildFromComplexType(xelem, *e.ComplexType)
		return xelem
	}

	if e.SimpleType != nil { // inline simple type
		b.buildFromSimpleType(xelem, *e.SimpleType)
		return xelem
	}

	return xelem
}

// buildFromComplexType takes an xmlElem and an xsdComplexType, containing
// XSD type information for xmlElem enrichment.
func (b builder) buildFromComplexType(xelem *xmlElem, t xsdComplexType) {
	if t.Sequence != nil { // Does the element have children?
		for _, e := range t.Sequence {
			xelem.Children = append(xelem.Children, b.buildFromElement(e))
		}
	}

	if t.Attributes != nil {
		b.buildFromAttributes(xelem, t.Attributes)
	}

	if t.ComplexContent != nil {
		b.buildFromComplexContent(xelem, *t.ComplexContent)
	}

	if t.SimpleContent != nil {
		b.buildFromSimpleContent(xelem, *t.SimpleContent)
	}
}

// buildFromSimpleType assumes restriction child and fetches the base value,
// assuming that value is of a XSD built-in data type.
func (b builder) buildFromSimpleType(xelem *xmlElem, t xsdSimpleType) {
	xelem.Type = b.findType(t.Restriction.Base).(string)
}

func (b builder) buildFromComplexContent(xelem *xmlElem, c xsdComplexContent) {
	if c.Extension != nil {
		b.buildFromExtension(xelem, c.Extension)
	}
}

// A simple content can refer to a text-only complex type
func (b builder) buildFromSimpleContent(xelem *xmlElem, c xsdSimpleContent) {
	if c.Extension != nil {
		b.buildFromExtension(xelem, c.Extension)
	}

	if c.Restriction != nil {
		b.buildFromRestriction(xelem, c.Restriction)
	}
}

// buildFromExtension extends an existing type, simple or complex, with a
// sequence.
func (b builder) buildFromExtension(xelem *xmlElem, e *xsdExtension) {
	switch t := b.findType(e.Base).(type) {
	case xsdComplexType:
		b.buildFromComplexType(xelem, t)
	case xsdSimpleType:
		b.buildFromSimpleType(xelem, t)
	default:
		xelem.Type = t.(string)
	}

	if e.Sequence != nil {
		for _, e := range e.Sequence {
			xelem.Children = append(xelem.Children, b.buildFromElement(e))
		}
	}

	if e.Attributes != nil {
		b.buildFromAttributes(xelem, e.Attributes)
	}
}

func (b builder) buildFromRestriction(xelem *xmlElem, r *xsdRestriction) {
	switch t := b.findType(r.Base).(type) {
	case xsdSimpleType:
		b.buildFromSimpleType(xelem, t)
	case xsdComplexType:
		b.buildFromComplexType(xelem, t)
	case xsdComplexContent:
		panic("Restriction on complex content is not implemented")
	default:
		panic("Unexpected base type to restriction")
	}
}

func typeFromXsdType(typ string) string {
	switch typ {
	case "boolean":
		typ = "bool"
	case "language", "dateTime", "Name", "token":
		typ = "string"
	case "long", "short", "integer", "int":
		typ = "int"
	case "decimal":
		typ = "float64"
	}
	return typ
}

func (b builder) buildFromAttributes(xelem *xmlElem, attrs []xsdAttribute) {
	for _, a := range attrs {
		attr := xmlAttrib{Name: a.Name}
		switch t := b.findType(a.Type).(type) {
		case xsdSimpleType:
			// Get type name from simpleType
			// If Restriction.Base is a simpleType or complexType, we panic
			attr.Type = b.findType(t.Restriction.Base).(string)
		case string:
			// If empty, then simpleType is present as content, but we ignore
			// that now
			attr.Type = t
		}
		xelem.Attribs = append(xelem.Attribs, attr)
	}
}

// findType takes a type name and checks if it is a registered XSD type
// (simple or complex), in which case that type is returned. If no such
// type can be found, the XSD specific primitive types are mapped to their
// Go correspondents. If no XSD type was found, the type name itself is
// returned.
func (b builder) findType(name string) interface{} {
	name = stripNamespace(name)
	if t, ok := b.complTypes[name]; ok {
		return t
	}
	if t, ok := b.simplTypes[name]; ok {
		return t
	}

	switch name {
	case "boolean":
		return "bool"
	case "language", "dateTime", "Name", "token":
		return "string"
	case "long", "short", "integer", "int":
		return "int"
	case "decimal":
		return "float64"
	default:
		return name
	}
}

func stripNamespace(name string) string {
	if s := strings.Split(name, ":"); len(s) > 1 {
		return s[len(s)-1]
	}
	return name
}
<|endoftext|>


<commit_before>package helps

import (
	"bytes"
	"encoding/xml"
	"io"
	"log"
)

func FormatXML(data []byte) string {
	formatted, err := FormatXMLSafe(data)
	if err != nil {
		log.Panic(err)
	}
	return formatted
}

// https://stackoverflow.com/a/27141132/605022
func FormatXMLSafe(data []byte) (string, error) {
	if len(data) > 3 && DumpFMT(data[:3]) == DumpFMT(utf8ByteOrderMark) {
		data = data[3:]
	}
	var buffer bytes.Buffer
	decoder := xml.NewDecoder(bytes.NewReader(data))
	encoder := xml.NewEncoder(&buffer)
	encoder.Indent("", "  ")
	for {
		token, err := decoder.Token()
		if err == io.EOF {
			encoder.Flush()
			return buffer.String(), nil
		}
		if err != nil {
			return "", err
		}
		err = encoder.EncodeToken(token)
		if err != nil {
			return "", err
		}
	}
}

var utf8ByteOrderMark = []byte{239, 187, 191}

func DumpXML(v interface{}) string {
	dump, err := DumpXMLSafe(v)
	if err != nil {
		log.Panic(err)
	}
	return dump
}

func DumpXMLSafe(v interface{}) (string, error) {
	dump, err := xml.MarshalIndent(v, "", "  ")
	if err != nil {
		return "", err
	}

	return string(dump), nil
}
<commit_msg>Extracted a type to manage state.<commit_after>package helps

import (
	"bytes"
	"encoding/xml"
	"io"
	"log"
)

func DumpXML(v interface{}) string {
	dump, err := DumpXMLSafe(v)
	if err != nil {
		log.Panic(err)
	}
	return dump
}

func DumpXMLSafe(v interface{}) (string, error) {
	dump, err := xml.MarshalIndent(v, "", "  ")
	if err != nil {
		return "", err
	}

	return string(dump), nil
}

func FormatXML(data []byte) string {
	formatted, err := FormatXMLSafe(data)
	if err != nil {
		log.Panic(err)
	}
	return formatted
}

func FormatXMLSafe(data []byte) (string, error) {
	return NewXMLFormatter(data).Format()
}

type XMLFormatter struct {
	buffer  *bytes.Buffer
	encoder *xml.Encoder
	decoder *xml.Decoder
}

func NewXMLFormatter(data []byte) *XMLFormatter {
	buffer := new(bytes.Buffer)
	return &XMLFormatter{
		buffer:  buffer,
		encoder: setupEncoder(buffer),
		decoder: setupDecoder(data),
	}
}

// https://stackoverflow.com/a/27141132/605022
func (this *XMLFormatter) Format() (string, error) {
	for {
		token, err := this.decoder.Token()
		if err == io.EOF {
			return this.finalResult()
		} else if err != nil {
			return "", err
		}

		err = this.encoder.EncodeToken(token)
		if err != nil {
			return "", err
		}
	}
}

func (this *XMLFormatter) finalResult() (string, error) {
	this.encoder.Flush()
	return this.buffer.String(), nil
}

func setupDecoder(data []byte) *xml.Decoder {
	return xml.NewDecoder(bytes.NewReader(trimByteOrderMark(data)))
}

func setupEncoder(buffer io.Writer) *xml.Encoder {
	encoder := xml.NewEncoder(buffer)
	encoder.Indent("", "  ")
	return encoder
}

func trimByteOrderMark(data []byte) []byte {
	return bytes.TrimLeft(data, string(utf8ByteOrderMark))
}

var utf8ByteOrderMark = []byte{239, 187, 191}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"

	_ "github.com/lib/pq"

	ydlconf "github.com/teo-mateo/ydl/config"
	"github.com/teo-mateo/ydl/handlers"
	"github.com/teo-mateo/ydl/util/cleanup"
)

func main() {
	fmt.Println("Starting YDL, port 8080.")
	fmt.Println("PG Conn: " + ydlconf.PgConnectionString())
	fmt.Println("Will convert to mp3")

	cleanup.StartCleanupRoutine()

	http.HandleFunc("/users", handlers.UsersHandler)
	http.HandleFunc("/ydl", handlers.QueueHandler)
	http.HandleFunc("/list/", handlers.ListHandler)
	http.HandleFunc("/list/json/", handlers.ListJSONPerUser)
	http.HandleFunc("/download", handlers.DownloadHandler)
	http.HandleFunc("/multidownload", handlers.MultiDownloadHandler)
	http.HandleFunc("/delete", handlers.DeleteHandler)
	http.HandleFunc("/static/", handlers.StaticFilesHandler)
	http.HandleFunc("/app", handlers.RootHandler)

	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>root handler to /<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"

	_ "github.com/lib/pq"

	ydlconf "github.com/teo-mateo/ydl/config"
	"github.com/teo-mateo/ydl/handlers"
	"github.com/teo-mateo/ydl/util/cleanup"
)

func main() {
	fmt.Println("Starting YDL, port 8080.")
	fmt.Println("PG Conn: " + ydlconf.PgConnectionString())
	fmt.Println("Will convert to mp3")

	cleanup.StartCleanupRoutine()

	http.HandleFunc("/", handlers.RootHandler)
	http.HandleFunc("/users", handlers.UsersHandler)
	http.HandleFunc("/ydl", handlers.QueueHandler)
	http.HandleFunc("/list/", handlers.ListHandler)
	http.HandleFunc("/list/json/", handlers.ListJSONPerUser)
	http.HandleFunc("/download", handlers.DownloadHandler)
	http.HandleFunc("/multidownload", handlers.MultiDownloadHandler)
	http.HandleFunc("/delete", handlers.DeleteHandler)
	http.HandleFunc("/static/", handlers.StaticFilesHandler)
	http.HandleFunc("/app", handlers.RootHandler)

	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>


<commit_before>package awspurge

import (
	"errors"
	"fmt"
	"net/http"
	"sync"
	"time"

	awsclient "github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/go-multierror"
)

type Config struct {
	Regions        []string      `toml:"regions" json:"regions"`
	RegionsExclude []string      `toml:"regions_exclude" json:"regions_exclude"`
	AccessKey      string        `toml:"access_key" json:"access_key"`
	SecretKey      string        `toml:"secret_key" json:"secret_key"`
	Timeout        time.Duration `toml:"timeout" json:"timeout"`
}

type resources struct {
	instances []*ec2.Instance
	volumes   []*ec2.Volume
}

type Purge struct {
	services *multiRegion

	// resources represents the current available resources per region. It's
	// populated by the Fetch() method.
	resources map[string]*resources
}

func New(conf *Config) (*Purge, error) {
	checkCfg := "Please check your configuration"

	if len(conf.Regions) == 0 {
		return nil, errors.New("AWS Regions are not set. " + checkCfg)
	}

	if conf.AccessKey == "" {
		return nil, errors.New("AWS Access Key is not set. " + checkCfg)
	}

	if conf.SecretKey == "" {
		return nil, errors.New("AWS Secret Key is not set. " + checkCfg)
	}

	if conf.Timeout == 0 {
		conf.Timeout = time.Second * 30
	}

	client := &http.Client{
		Transport: &http.Transport{TLSHandshakeTimeout: conf.Timeout},
		Timeout:   conf.Timeout,
	}

	creds := credentials.NewStaticCredentials(conf.AccessKey, conf.SecretKey, "")
	awsCfg := &awsclient.Config{
		Credentials: creds,
		HTTPClient:  client,
		Logger:      awsclient.NewDefaultLogger(),
	}

	m := newMultiRegion(awsCfg, filterRegions(conf.Regions, conf.RegionsExclude))
	return &Purge{
		services:  m,
		resources: make(map[string]*resources, 0),
	}, nil
}

func (p *Purge) Do() error {
	if err := p.Fetch(); err != nil {
		return err
	}

	if err := p.Print(); err != nil {
		return err
	}

	return nil
}

// Print prints all fetched resources
func (p *Purge) Print() error {
	for region, resources := range p.resources {
		fmt.Printf("[%s] found '%d' instances\n", region, len(resources.instances))
		fmt.Printf("[%s] found '%d' volumes\n", region, len(resources.volumes))
	}
	return nil
}

// Fetch fetches all given resources and stores them internally. To print them
// use the Print() method
func (p *Purge) Fetch() error {
	allInstances, err := p.DescribeInstances()
	if err != nil {
		return err
	}

	allVolumes, err := p.DescribeVolumes()
	if err != nil {
		return err
	}

	for _, region := range allRegions {
		p.resources[region] = &resources{
			instances: allInstances[region],
			volumes:   allVolumes[region],
		}
	}

	return nil
}

// DescribeVolumes returns all volumes per region.
func (p *Purge) DescribeVolumes() (map[string][]*ec2.Volume, error) {
	var (
		wg sync.WaitGroup
		mu sync.Mutex

		multiErrors error
	)

	output := make(map[string][]*ec2.Volume)

	for r, s := range p.services.regions {
		wg.Add(1)

		go func(region string, svc *ec2.EC2) {
			resp, err := svc.DescribeVolumes(nil)
			if err != nil {
				mu.Lock()
				multiErrors = multierror.Append(multiErrors, err)
				mu.Unlock()
				wg.Done()
				return
			}

			if resp.Volumes != nil && len(resp.Volumes) != 0 {
				mu.Lock()
				output[region] = resp.Volumes
				mu.Unlock()
			}

			wg.Done()
		}(r, s)
	}

	wg.Wait()

	return output, multiErrors
}

// DescribeInstances returns all instances per region.
func (p *Purge) DescribeInstances() (map[string][]*ec2.Instance, error) {
	var (
		wg sync.WaitGroup
		mu sync.Mutex

		multiErrors error
	)

	output := make(map[string][]*ec2.Instance)

	for r, s := range p.services.regions {
		wg.Add(1)

		go func(region string, svc *ec2.EC2) {
			resp, err := svc.DescribeInstances(nil)
			if err != nil {
				mu.Lock()
				multiErrors = multierror.Append(multiErrors, err)
				mu.Unlock()
				wg.Done()
				return
			}

			if resp.Reservations != nil {
				instances := make([]*ec2.Instance, 0)

				for _, reserv := range resp.Reservations {
					if len(reserv.Instances) != 0 {
						instances = append(instances, reserv.Instances...)
					}
				}

				mu.Lock()
				output[region] = instances
				mu.Unlock()
			}

			wg.Done()
		}(r, s)
	}

	wg.Wait()

	return output, multiErrors
}
<commit_msg>awspurge: wrap multi region fetching<commit_after>package awspurge

import (
	"errors"
	"fmt"
	"net/http"
	"sync"
	"time"

	awsclient "github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/go-multierror"
)

type Config struct {
	Regions        []string      `toml:"regions" json:"regions"`
	RegionsExclude []string      `toml:"regions_exclude" json:"regions_exclude"`
	AccessKey      string        `toml:"access_key" json:"access_key"`
	SecretKey      string        `toml:"secret_key" json:"secret_key"`
	Timeout        time.Duration `toml:"timeout" json:"timeout"`
}

type resources struct {
	instances []*ec2.Instance
	volumes   []*ec2.Volume
}

type Purge struct {
	services *multiRegion

	// resources represents the current available resources per region. It's
	// populated by the Fetch() method.
	resources map[string]*resources
}

func New(conf *Config) (*Purge, error) {
	checkCfg := "Please check your configuration"

	if len(conf.Regions) == 0 {
		return nil, errors.New("AWS Regions are not set. " + checkCfg)
	}

	if conf.AccessKey == "" {
		return nil, errors.New("AWS Access Key is not set. " + checkCfg)
	}

	if conf.SecretKey == "" {
		return nil, errors.New("AWS Secret Key is not set. " + checkCfg)
	}

	if conf.Timeout == 0 {
		conf.Timeout = time.Second * 30
	}

	client := &http.Client{
		Transport: &http.Transport{TLSHandshakeTimeout: conf.Timeout},
		Timeout:   conf.Timeout,
	}

	creds := credentials.NewStaticCredentials(conf.AccessKey, conf.SecretKey, "")
	awsCfg := &awsclient.Config{
		Credentials: creds,
		HTTPClient:  client,
		Logger:      awsclient.NewDefaultLogger(),
	}

	m := newMultiRegion(awsCfg, filterRegions(conf.Regions, conf.RegionsExclude))
	return &Purge{
		services:  m,
		resources: make(map[string]*resources, 0),
	}, nil
}

func (p *Purge) Do() error {
	if err := p.Fetch(); err != nil {
		return err
	}

	if err := p.Print(); err != nil {
		return err
	}

	return nil
}

// Print prints all fetched resources
func (p *Purge) Print() error {
	for region, resources := range p.resources {
		fmt.Printf("[%s] found '%d' instances\n", region, len(resources.instances))
		fmt.Printf("[%s] found '%d' volumes\n", region, len(resources.volumes))
	}
	return nil
}

// Fetch fetches all given resources and stores them internally. To print them
// use the Print() method
func (p *Purge) Fetch() error {
	allInstances, err := p.DescribeInstances()
	if err != nil {
		return err
	}

	allVolumes, err := p.DescribeVolumes()
	if err != nil {
		return err
	}

	for _, region := range allRegions {
		p.resources[region] = &resources{
			instances: allInstances[region],
			volumes:   allVolumes[region],
		}
	}

	return nil
}

func (p *Purge) describeResources(fn func(*ec2.EC2) (interface{}, error)) (map[string]interface{}, error) {
	var (
		wg sync.WaitGroup
		mu sync.Mutex

		multiErrors error
	)

	output := make(map[string]interface{})

	for r, s := range p.services.regions {
		wg.Add(1)

		go func(region string, svc *ec2.EC2) {
			out, err := fn(svc)
			if err != nil {
				mu.Lock()
				multiErrors = multierror.Append(multiErrors, err)
				mu.Unlock()
				wg.Done()
				return
			}

			mu.Lock()
			output[region] = out
			mu.Unlock()

			wg.Done()
		}(r, s)
	}

	wg.Wait()

	return output, multiErrors
}

// DescribeVolumes returns all volumes per region.
func (p *Purge) DescribeVolumes() (map[string][]*ec2.Volume, error) {
	describeVolumes := func(svc *ec2.EC2) (interface{}, error) {
		resp, err := svc.DescribeVolumes(nil)
		if err != nil {
			return nil, err
		}

		return resp.Volumes, nil
	}

	out, err := p.describeResources(describeVolumes)
	if err != nil {
		return nil, err
	}

	volumes := make(map[string][]*ec2.Volume)
	for region, v := range out {
		vols, ok := v.([]*ec2.Volume)
		if !ok {
			continue
		}
		volumes[region] = vols
	}

	return volumes, nil
}

// DescribeInstances returns all instances per region.
func (p *Purge) DescribeInstances() (map[string][]*ec2.Instance, error) {
	describeInstances := func(svc *ec2.EC2) (interface{}, error) {
		resp, err := svc.DescribeInstances(nil)
		if err != nil {
			return nil, err
		}

		instances := make([]*ec2.Instance, 0)
		if resp.Reservations != nil {
			for _, reserv := range resp.Reservations {
				if len(reserv.Instances) != 0 {
					instances = append(instances, reserv.Instances...)
				}
			}
		}

		return instances, nil
	}

	out, err := p.describeResources(describeInstances)
	if err != nil {
		return nil, err
	}

	instances := make(map[string][]*ec2.Instance)
	for region, i := range out {
		ins, ok := i.([]*ec2.Instance)
		if !ok {
			continue
		}
		instances[region] = ins
	}

	return instances, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package backups

import (
	"compress/gzip"
	"crypto/sha1"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/errors"
	"github.com/juju/utils/hash"
	"github.com/juju/utils/tar"

	"github.com/juju/juju/state/backups/archive"
)

// TODO(ericsnow) One concern is files that get out of date by the time
// backup finishes running.  This is particularly a problem with log
// files.

const (
	tempPrefix   = "jujuBackup-"
	tempFilename = "juju-backup.tar.gz"
)

type dumper interface {
	Dump(dumpDir string) error
}

type createArgs struct {
	filesToBackUp []string
	db            dumper
}

type createResult struct {
	archiveFile io.ReadCloser
	size        int64
	checksum    string
}

// create builds a new backup archive file and returns it.  It also
// updates the metadata with the file info.
func create(args *createArgs) (*createResult, error) {
	// Prepare the backup builder.
	builder, err := newBuilder(args.filesToBackUp, args.db)
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer builder.cleanUp()

	// Build the backup.
	if err := builder.buildAll(); err != nil {
		return nil, errors.Trace(err)
	}

	// Get the result.
	result, err := builder.result()
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Return the result.  Note that the entire build workspace will be
	// deleted at the end of this function.  This includes the backup
	// archive file we built.  However, the handle to that file in the
	// result will still be open and readable.
	// If we ever support state machines on Windows, this will need to
	// change (you can't delete open files on Windows).
	return result, nil
}

// builder exposes the machinery for creating a backup of juju's state.
type builder struct {
	// archive is the backups archive summary.
	archive *archive.Archive
	// checksum is the checksum of the archive file.
	checksum string
	// filesToBackUp is the paths to every file to include in the archive.
	filesToBackUp []string
	// db is the wrapper around the DB dump command and args.
	db dumper
	// archiveFile is the backup archive file.
	archiveFile io.WriteCloser
	// bundleFile is the inner archive file containing all the juju
	// state-related files gathered during backup.
	bundleFile io.WriteCloser
}

// newBuilder returns a new backup archive builder.  It creates the temp
// directories which backup uses as its staging area while building the
// archive.  It also creates the archive
// (temp root, tarball root, DB dumpdir), along with any error.
func newBuilder(filesToBackUp []string, db dumper) (*builder, error) {
	b := builder{
		filesToBackUp: filesToBackUp,
		db:            db,
	}

	// Create the backups workspace root directory.
	rootDir, err := ioutil.TempDir("", tempPrefix)
	if err != nil {
		return nil, errors.Annotate(err, "while making backups workspace")
	}
	filename := filepath.Join(rootDir, tempFilename)
	b.archive = &archive.Archive{filename, rootDir}

	// Create all the direcories we need.  We go with user-only
	// permissions on principle; the directories are short-lived so in
	// practice it shouldn't matter much.
	err = os.MkdirAll(b.archive.DBDumpDir(), 0700)
	if err != nil {
		b.cleanUp()
		return nil, errors.Annotate(err, "while creating temp directories")
	}

	// Create the archive files.  We do so here to fail as early as
	// possible.
	b.archiveFile, err = os.Create(filename)
	if err != nil {
		b.cleanUp()
		return nil, errors.Annotate(err, "while creating archive file")
	}

	b.bundleFile, err = os.Create(b.archive.FilesBundle())
	if err != nil {
		b.cleanUp()
		return nil, errors.Annotate(err, `while creating bundle file`)
	}

	return &b, nil
}

func (b *builder) closeArchiveFile() error {
	if b.archiveFile == nil {
		return nil
	}

	if err := b.archiveFile.Close(); err != nil {
		return errors.Annotate(err, "while closing archive file")
	}

	b.archiveFile = nil
	return nil
}

func (b *builder) closeBundleFile() error {
	if b.bundleFile == nil {
		return nil
	}

	if err := b.bundleFile.Close(); err != nil {
		return errors.Annotate(err, "while closing bundle file")
	}

	b.bundleFile = nil
	return nil
}

func (b *builder) removeRootDir() error {
	if b.archive == nil || b.archive.UnpackedRootDir == "" {
		return nil
	}

	if err := os.RemoveAll(b.archive.UnpackedRootDir); err != nil {
		return errors.Annotate(err, "while removing backups temp dir")
	}

	return nil
}

func (b *builder) cleanUp() error {
	var failed int

	funcs := [](func() error){
		b.closeBundleFile,
		b.closeArchiveFile,
		b.removeRootDir,
	}
	for _, cleanupFunc := range funcs {
		if err := cleanupFunc(); err != nil {
			logger.Errorf(err.Error())
			failed++
		}
	}

	if failed > 0 {
		return errors.Errorf("%d errors during cleanup (see logs)", failed)
	}
	return nil
}

func (b *builder) buildFilesBundle() error {
	logger.Infof("dumping juju state-related files")
	if b.filesToBackUp == nil {
		logger.Infof("nothing to do")
		return nil
	}
	if b.bundleFile == nil {
		return errors.New("missing bundleFile")
	}

	stripPrefix := string(os.PathSeparator)
	_, err := tar.TarFiles(b.filesToBackUp, b.bundleFile, stripPrefix)
	if err != nil {
		return errors.Annotate(err, "while bundling state-critical files")
	}

	return nil
}

func (b *builder) buildDBDump() error {
	logger.Infof("dumping database")
	if b.db == nil {
		logger.Infof("nothing to do")
		return nil
	}

	dumpDir := b.archive.DBDumpDir()
	if err := b.db.Dump(dumpDir); err != nil {
		return errors.Annotate(err, "while dumping juju state database")
	}

	return nil
}

func (b *builder) buildArchive(outFile io.Writer) error {
	tarball := gzip.NewWriter(outFile)
	defer tarball.Close()

	// We add a trailing slash (or whatever) to root so that everything
	// in the path up to and including that slash is stripped off when
	// each file is added to the tar file.
	stripPrefix := b.archive.UnpackedRootDir + string(os.PathSeparator)
	filenames := []string{b.archive.ContentDir()}
	if _, err := tar.TarFiles(filenames, tarball, stripPrefix); err != nil {
		return errors.Annotate(err, "while bundling final archive")
	}

	return nil
}

func (b *builder) buildArchiveAndChecksum() error {
	logger.Infof("building archive file (%s)", b.archive.Filename)
	if b.archiveFile == nil {
		return errors.New("missing archiveFile")
	}

	// Build the tarball, writing out to both the archive file and a
	// SHA1 hash.  The hash will correspond to the gzipped file rather
	// than to the uncompressed contents of the tarball.  This is so
	// that users can compare the published checksum against the
	// checksum of the file without having to decompress it first.
	hasher := hash.NewHashingWriter(b.archiveFile, sha1.New())
	if err := b.buildArchive(hasher); err != nil {
		return errors.Trace(err)
	}

	// Save the SHA1 checksum.
	// Gzip writers may buffer what they're writing so we must call
	// Close() on the writer *before* getting the checksum from the
	// hasher.
	b.checksum = hasher.Base64Sum()

	return nil
}

func (b *builder) buildAll() error {
	// Dump the files.
	if err := b.buildFilesBundle(); err != nil {
		return errors.Trace(err)
	}

	// Dump the database.
	if err := b.buildDBDump(); err != nil {
		return errors.Trace(err)
	}

	// Bundle it all into a tarball.
	if err := b.buildArchiveAndChecksum(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (b *builder) result() (*createResult, error) {
	// Open the file in read-only mode.
	file, err := os.Open(b.archive.Filename)
	if err != nil {
		return nil, errors.Annotate(err, "while opening archive file")
	}

	// Get the size.
	stat, err := file.Stat()
	if err != nil {
		return nil, errors.Annotate(err, "while reading archive file info")
	}
	size := stat.Size()

	// Get the checksum.
	checksum := b.checksum

	// Return the result.
	result := createResult{
		archiveFile: file,
		size:        size,
		checksum:    checksum,
	}
	return &result, nil
}
<commit_msg>Use defer for cleaning up.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package backups

import (
	"compress/gzip"
	"crypto/sha1"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/errors"
	"github.com/juju/utils/hash"
	"github.com/juju/utils/tar"

	"github.com/juju/juju/state/backups/archive"
)

// TODO(ericsnow) One concern is files that get out of date by the time
// backup finishes running.  This is particularly a problem with log
// files.

const (
	tempPrefix   = "jujuBackup-"
	tempFilename = "juju-backup.tar.gz"
)

type dumper interface {
	Dump(dumpDir string) error
}

type createArgs struct {
	filesToBackUp []string
	db            dumper
}

type createResult struct {
	archiveFile io.ReadCloser
	size        int64
	checksum    string
}

// create builds a new backup archive file and returns it.  It also
// updates the metadata with the file info.
func create(args *createArgs) (*createResult, error) {
	// Prepare the backup builder.
	builder, err := newBuilder(args.filesToBackUp, args.db)
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer builder.cleanUp()

	// Build the backup.
	if err := builder.buildAll(); err != nil {
		return nil, errors.Trace(err)
	}

	// Get the result.
	result, err := builder.result()
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Return the result.  Note that the entire build workspace will be
	// deleted at the end of this function.  This includes the backup
	// archive file we built.  However, the handle to that file in the
	// result will still be open and readable.
	// If we ever support state machines on Windows, this will need to
	// change (you can't delete open files on Windows).
	return result, nil
}

// builder exposes the machinery for creating a backup of juju's state.
type builder struct {
	// archive is the backups archive summary.
	archive *archive.Archive
	// checksum is the checksum of the archive file.
	checksum string
	// filesToBackUp is the paths to every file to include in the archive.
	filesToBackUp []string
	// db is the wrapper around the DB dump command and args.
	db dumper
	// archiveFile is the backup archive file.
	archiveFile io.WriteCloser
	// bundleFile is the inner archive file containing all the juju
	// state-related files gathered during backup.
	bundleFile io.WriteCloser
}

// newBuilder returns a new backup archive builder.  It creates the temp
// directories which backup uses as its staging area while building the
// archive.  It also creates the archive
// (temp root, tarball root, DB dumpdir), along with any error.
func newBuilder(filesToBackUp []string, db dumper) (*builder, error) {
	var err error

	b := builder{
		filesToBackUp: filesToBackUp,
		db:            db,
	}
	defer func() {
		if err != nil {
			if err := b.cleanUp(); err != nil {
				logger.Errorf("error while cleaning up: %v", err)
			}
		}
	}()

	// Create the backups workspace root directory.
	rootDir, err := ioutil.TempDir("", tempPrefix)
	if err != nil {
		return nil, errors.Annotate(err, "while making backups workspace")
	}
	filename := filepath.Join(rootDir, tempFilename)
	b.archive = &archive.Archive{filename, rootDir}

	// Create all the direcories we need.  We go with user-only
	// permissions on principle; the directories are short-lived so in
	// practice it shouldn't matter much.
	err = os.MkdirAll(b.archive.DBDumpDir(), 0700)
	if err != nil {
		return nil, errors.Annotate(err, "while creating temp directories")
	}

	// Create the archive files.  We do so here to fail as early as
	// possible.
	b.archiveFile, err = os.Create(filename)
	if err != nil {
		return nil, errors.Annotate(err, "while creating archive file")
	}

	b.bundleFile, err = os.Create(b.archive.FilesBundle())
	if err != nil {
		return nil, errors.Annotate(err, `while creating bundle file`)
	}

	return &b, nil
}

func (b *builder) closeArchiveFile() error {
	if b.archiveFile == nil {
		return nil
	}

	if err := b.archiveFile.Close(); err != nil {
		return errors.Annotate(err, "while closing archive file")
	}

	b.archiveFile = nil
	return nil
}

func (b *builder) closeBundleFile() error {
	if b.bundleFile == nil {
		return nil
	}

	if err := b.bundleFile.Close(); err != nil {
		return errors.Annotate(err, "while closing bundle file")
	}

	b.bundleFile = nil
	return nil
}

func (b *builder) removeRootDir() error {
	if b.archive == nil || b.archive.UnpackedRootDir == "" {
		return nil
	}

	if err := os.RemoveAll(b.archive.UnpackedRootDir); err != nil {
		return errors.Annotate(err, "while removing backups temp dir")
	}

	return nil
}

func (b *builder) cleanUp() error {
	var failed int

	funcs := [](func() error){
		b.closeBundleFile,
		b.closeArchiveFile,
		b.removeRootDir,
	}
	for _, cleanupFunc := range funcs {
		if err := cleanupFunc(); err != nil {
			logger.Errorf(err.Error())
			failed++
		}
	}

	if failed > 0 {
		return errors.Errorf("%d errors during cleanup (see logs)", failed)
	}
	return nil
}

func (b *builder) buildFilesBundle() error {
	logger.Infof("dumping juju state-related files")
	if b.filesToBackUp == nil {
		logger.Infof("nothing to do")
		return nil
	}
	if b.bundleFile == nil {
		return errors.New("missing bundleFile")
	}

	stripPrefix := string(os.PathSeparator)
	_, err := tar.TarFiles(b.filesToBackUp, b.bundleFile, stripPrefix)
	if err != nil {
		return errors.Annotate(err, "while bundling state-critical files")
	}

	return nil
}

func (b *builder) buildDBDump() error {
	logger.Infof("dumping database")
	if b.db == nil {
		logger.Infof("nothing to do")
		return nil
	}

	dumpDir := b.archive.DBDumpDir()
	if err := b.db.Dump(dumpDir); err != nil {
		return errors.Annotate(err, "while dumping juju state database")
	}

	return nil
}

func (b *builder) buildArchive(outFile io.Writer) error {
	tarball := gzip.NewWriter(outFile)
	defer tarball.Close()

	// We add a trailing slash (or whatever) to root so that everything
	// in the path up to and including that slash is stripped off when
	// each file is added to the tar file.
	stripPrefix := b.archive.UnpackedRootDir + string(os.PathSeparator)
	filenames := []string{b.archive.ContentDir()}
	if _, err := tar.TarFiles(filenames, tarball, stripPrefix); err != nil {
		return errors.Annotate(err, "while bundling final archive")
	}

	return nil
}

func (b *builder) buildArchiveAndChecksum() error {
	logger.Infof("building archive file (%s)", b.archive.Filename)
	if b.archiveFile == nil {
		return errors.New("missing archiveFile")
	}

	// Build the tarball, writing out to both the archive file and a
	// SHA1 hash.  The hash will correspond to the gzipped file rather
	// than to the uncompressed contents of the tarball.  This is so
	// that users can compare the published checksum against the
	// checksum of the file without having to decompress it first.
	hasher := hash.NewHashingWriter(b.archiveFile, sha1.New())
	if err := b.buildArchive(hasher); err != nil {
		return errors.Trace(err)
	}

	// Save the SHA1 checksum.
	// Gzip writers may buffer what they're writing so we must call
	// Close() on the writer *before* getting the checksum from the
	// hasher.
	b.checksum = hasher.Base64Sum()

	return nil
}

func (b *builder) buildAll() error {
	// Dump the files.
	if err := b.buildFilesBundle(); err != nil {
		return errors.Trace(err)
	}

	// Dump the database.
	if err := b.buildDBDump(); err != nil {
		return errors.Trace(err)
	}

	// Bundle it all into a tarball.
	if err := b.buildArchiveAndChecksum(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (b *builder) result() (*createResult, error) {
	// Open the file in read-only mode.
	file, err := os.Open(b.archive.Filename)
	if err != nil {
		return nil, errors.Annotate(err, "while opening archive file")
	}

	// Get the size.
	stat, err := file.Stat()
	if err != nil {
		return nil, errors.Annotate(err, "while reading archive file info")
	}
	size := stat.Size()

	// Get the checksum.
	checksum := b.checksum

	// Return the result.
	result := createResult{
		archiveFile: file,
		size:        size,
		checksum:    checksum,
	}
	return &result, nil
}
<|endoftext|>


<commit_before>// Copyright 2011 Gary Burd
//
// 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 redis

import (
	"io"
	"reflect"
	"testing"
	"time"
)

type poolTestConn struct {
	d   *poolDialer
	err error
	Conn
}

func (c *poolTestConn) Close() error { c.d.open -= 1; return nil }
func (c *poolTestConn) Err() error   { return c.err }

func (c *poolTestConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
	if commandName == "ERR" {
		c.err = args[0].(error)
	}
	if commandName != "" {
		c.d.commands = append(c.d.commands, commandName)
	}
	return c.Conn.Do(commandName, args...)
}

func (c *poolTestConn) Send(commandName string, args ...interface{}) error {
	c.d.commands = append(c.d.commands, commandName)
	return c.Conn.Send(commandName, args...)
}

type poolDialer struct {
	t            *testing.T
	dialed, open int
	commands     []string
}

func (d *poolDialer) dial() (Conn, error) {
	d.open += 1
	d.dialed += 1
	c, err := DialTestDB()
	if err != nil {
		return nil, err
	}
	return &poolTestConn{d: d, Conn: c}, nil
}

func (d *poolDialer) check(message string, p *Pool, dialed, open int) {
	if d.dialed != dialed {
		d.t.Errorf("%s: dialed=%d, want %d", message, d.dialed, dialed)
	}
	if d.open != open {
		d.t.Errorf("%s: open=%d, want %d", message, d.open, open)
	}
	if active := p.ActiveCount(); active != open {
		d.t.Errorf("%s: active=%d, want %d", message, active, open)
	}
}

func TestPoolReuse(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}

	for i := 0; i < 10; i++ {
		c1 := p.Get()
		c1.Do("PING")
		c2 := p.Get()
		c2.Do("PING")
		c1.Close()
		c2.Close()
	}

	d.check("before close", p, 2, 2)
	p.Close()
	d.check("after close", p, 2, 0)
}

func TestPoolMaxIdle(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}
	for i := 0; i < 10; i++ {
		c1 := p.Get()
		c1.Do("PING")
		c2 := p.Get()
		c2.Do("PING")
		c3 := p.Get()
		c3.Do("PING")
		c1.Close()
		c2.Close()
		c3.Close()
	}
	d.check("before close", p, 12, 2)
	p.Close()
	d.check("after close", p, 12, 0)
}

func TestPoolError(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}

	c := p.Get()
	c.Do("ERR", io.EOF)
	if c.Err() == nil {
		t.Errorf("expected c.Err() != nil")
	}
	c.Close()

	c = p.Get()
	c.Do("ERR", io.EOF)
	c.Close()

	d.check(".", p, 2, 0)
}

func TestPoolClose(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}

	c1 := p.Get()
	c1.Do("PING")
	c2 := p.Get()
	c2.Do("PING")
	c3 := p.Get()
	c3.Do("PING")

	c1.Close()
	if _, err := c1.Do("PING"); err == nil {
		t.Errorf("expected error after connection closed")
	}

	c2.Close()
	c2.Close()

	p.Close()

	d.check("after pool close", p, 3, 1)

	if _, err := c1.Do("PING"); err == nil {
		t.Errorf("expected error after connection and pool closed")
	}

	c3.Close()

	d.check("after conn close", p, 3, 0)

	c1 = p.Get()
	if _, err := c1.Do("PING"); err == nil {
		t.Errorf("expected error after pool closed")
	}
}

func TestPoolTimeout(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:     2,
		IdleTimeout: 300 * time.Second,
		Dial:        d.dial,
	}

	now := time.Now()
	nowFunc = func() time.Time { return now }
	defer func() { nowFunc = time.Now }()

	c := p.Get()
	c.Do("PING")
	c.Close()

	d.check("1", p, 1, 1)

	now = now.Add(p.IdleTimeout)

	c = p.Get()
	c.Do("PING")
	c.Close()

	d.check("2", p, 2, 1)

	p.Close()
}

func TestPoolConcurrenSendReceive(t *testing.T) {
	p := &Pool{
		Dial: DialTestDB,
	}
	c := p.Get()
	done := make(chan error, 1)
	go func() {
		_, err := c.Receive()
		done <- err
	}()
	c.Send("PING")
	c.Flush()
	err := <-done
	if err != nil {
		t.Fatalf("Receive() returned error %v", err)
	}
	_, err = c.Do("")
	if err != nil {
		t.Fatalf("Do() returned error %v", err)
	}
	c.Close()
	p.Close()
}

func TestPoolBorrowCheck(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:      2,
		Dial:         d.dial,
		TestOnBorrow: func(Conn, time.Time) error { return Error("BLAH") },
	}

	for i := 0; i < 10; i++ {
		c := p.Get()
		c.Do("PING")
		c.Close()
	}
	d.check("1", p, 10, 1)
	p.Close()
}

func TestPoolMaxActive(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}
	c1 := p.Get()
	c1.Do("PING")
	c2 := p.Get()
	c2.Do("PING")

	d.check("1", p, 2, 2)

	c3 := p.Get()
	if _, err := c3.Do("PING"); err != ErrPoolExhausted {
		t.Errorf("expected pool exhausted")
	}

	c3.Close()
	d.check("2", p, 2, 2)
	c2.Close()
	d.check("3", p, 2, 2)

	c3 = p.Get()
	if _, err := c3.Do("PING"); err != nil {
		t.Errorf("expected good channel, err=%v", err)
	}
	c3.Close()

	d.check("4", p, 2, 2)
	p.Close()
}

func TestPoolMonitorCleanup(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}
	c := p.Get()
	c.Send("MONITOR")
	c.Close()

	d.check("", p, 1, 0)
	p.Close()
}

func TestPoolPubSubCleanup(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}

	c := p.Get()
	c.Send("SUBSCRIBE", "x")
	c.Close()

	want := []string{"SUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", "ECHO"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Send("PSUBSCRIBE", "x*")
	c.Close()

	want = []string{"PSUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", "ECHO"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	p.Close()
}

func TestPoolTransactionCleanup(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}

	c := p.Get()
	c.Do("WATCH", "key")
	c.Do("PING")
	c.Close()

	want := []string{"WATCH", "PING", "UNWATCH"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("UNWATCH")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "UNWATCH", "PING"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("MULTI")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "MULTI", "PING", "DISCARD"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("MULTI")
	c.Do("DISCARD")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "MULTI", "DISCARD", "PING"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("MULTI")
	c.Do("EXEC")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "MULTI", "EXEC", "PING"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	p.Close()
}

func BenchmarkPoolGet(b *testing.B) {
	b.StopTimer()
	p := Pool{Dial: DialTestDB, MaxIdle: 2}
	c := p.Get()
	if err := c.Err(); err != nil {
		b.Fatal(err)
	}
	c.Close()
	defer p.Close()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		c = p.Get()
		c.Close()
	}
}

func BenchmarkPoolGetErr(b *testing.B) {
	b.StopTimer()
	p := Pool{Dial: DialTestDB, MaxIdle: 2}
	c := p.Get()
	if err := c.Err(); err != nil {
		b.Fatal(err)
	}
	c.Close()
	defer p.Close()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		c = p.Get()
		if err := c.Err(); err != nil {
			b.Fatal(err)
		}
		c.Close()
	}
}

func BenchmarkPoolGetPing(b *testing.B) {
	b.StopTimer()
	p := Pool{Dial: DialTestDB, MaxIdle: 2}
	c := p.Get()
	if err := c.Err(); err != nil {
		b.Fatal(err)
	}
	c.Close()
	defer p.Close()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		c = p.Get()
		if _, err := c.Do("PING"); err != nil {
			b.Fatal(err)
		}
		c.Close()
	}
}
<commit_msg>Add concurrency benchmarks<commit_after>// Copyright 2011 Gary Burd
//
// 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 redis

import (
	"io"
	"net/textproto"
	"reflect"
	"sync"
	"testing"
	"time"
)

type poolTestConn struct {
	d   *poolDialer
	err error
	Conn
}

func (c *poolTestConn) Close() error { c.d.open -= 1; return nil }
func (c *poolTestConn) Err() error   { return c.err }

func (c *poolTestConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
	if commandName == "ERR" {
		c.err = args[0].(error)
	}
	if commandName != "" {
		c.d.commands = append(c.d.commands, commandName)
	}
	return c.Conn.Do(commandName, args...)
}

func (c *poolTestConn) Send(commandName string, args ...interface{}) error {
	c.d.commands = append(c.d.commands, commandName)
	return c.Conn.Send(commandName, args...)
}

type poolDialer struct {
	t            *testing.T
	dialed, open int
	commands     []string
}

func (d *poolDialer) dial() (Conn, error) {
	d.open += 1
	d.dialed += 1
	c, err := DialTestDB()
	if err != nil {
		return nil, err
	}
	return &poolTestConn{d: d, Conn: c}, nil
}

func (d *poolDialer) check(message string, p *Pool, dialed, open int) {
	if d.dialed != dialed {
		d.t.Errorf("%s: dialed=%d, want %d", message, d.dialed, dialed)
	}
	if d.open != open {
		d.t.Errorf("%s: open=%d, want %d", message, d.open, open)
	}
	if active := p.ActiveCount(); active != open {
		d.t.Errorf("%s: active=%d, want %d", message, active, open)
	}
}

func TestPoolReuse(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}

	for i := 0; i < 10; i++ {
		c1 := p.Get()
		c1.Do("PING")
		c2 := p.Get()
		c2.Do("PING")
		c1.Close()
		c2.Close()
	}

	d.check("before close", p, 2, 2)
	p.Close()
	d.check("after close", p, 2, 0)
}

func TestPoolMaxIdle(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}
	for i := 0; i < 10; i++ {
		c1 := p.Get()
		c1.Do("PING")
		c2 := p.Get()
		c2.Do("PING")
		c3 := p.Get()
		c3.Do("PING")
		c1.Close()
		c2.Close()
		c3.Close()
	}
	d.check("before close", p, 12, 2)
	p.Close()
	d.check("after close", p, 12, 0)
}

func TestPoolError(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}

	c := p.Get()
	c.Do("ERR", io.EOF)
	if c.Err() == nil {
		t.Errorf("expected c.Err() != nil")
	}
	c.Close()

	c = p.Get()
	c.Do("ERR", io.EOF)
	c.Close()

	d.check(".", p, 2, 0)
}

func TestPoolClose(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle: 2,
		Dial:    d.dial,
	}

	c1 := p.Get()
	c1.Do("PING")
	c2 := p.Get()
	c2.Do("PING")
	c3 := p.Get()
	c3.Do("PING")

	c1.Close()
	if _, err := c1.Do("PING"); err == nil {
		t.Errorf("expected error after connection closed")
	}

	c2.Close()
	c2.Close()

	p.Close()

	d.check("after pool close", p, 3, 1)

	if _, err := c1.Do("PING"); err == nil {
		t.Errorf("expected error after connection and pool closed")
	}

	c3.Close()

	d.check("after conn close", p, 3, 0)

	c1 = p.Get()
	if _, err := c1.Do("PING"); err == nil {
		t.Errorf("expected error after pool closed")
	}
}

func TestPoolTimeout(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:     2,
		IdleTimeout: 300 * time.Second,
		Dial:        d.dial,
	}

	now := time.Now()
	nowFunc = func() time.Time { return now }
	defer func() { nowFunc = time.Now }()

	c := p.Get()
	c.Do("PING")
	c.Close()

	d.check("1", p, 1, 1)

	now = now.Add(p.IdleTimeout)

	c = p.Get()
	c.Do("PING")
	c.Close()

	d.check("2", p, 2, 1)

	p.Close()
}

func TestPoolConcurrenSendReceive(t *testing.T) {
	p := &Pool{
		Dial: DialTestDB,
	}
	c := p.Get()
	done := make(chan error, 1)
	go func() {
		_, err := c.Receive()
		done <- err
	}()
	c.Send("PING")
	c.Flush()
	err := <-done
	if err != nil {
		t.Fatalf("Receive() returned error %v", err)
	}
	_, err = c.Do("")
	if err != nil {
		t.Fatalf("Do() returned error %v", err)
	}
	c.Close()
	p.Close()
}

func TestPoolBorrowCheck(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:      2,
		Dial:         d.dial,
		TestOnBorrow: func(Conn, time.Time) error { return Error("BLAH") },
	}

	for i := 0; i < 10; i++ {
		c := p.Get()
		c.Do("PING")
		c.Close()
	}
	d.check("1", p, 10, 1)
	p.Close()
}

func TestPoolMaxActive(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}
	c1 := p.Get()
	c1.Do("PING")
	c2 := p.Get()
	c2.Do("PING")

	d.check("1", p, 2, 2)

	c3 := p.Get()
	if _, err := c3.Do("PING"); err != ErrPoolExhausted {
		t.Errorf("expected pool exhausted")
	}

	c3.Close()
	d.check("2", p, 2, 2)
	c2.Close()
	d.check("3", p, 2, 2)

	c3 = p.Get()
	if _, err := c3.Do("PING"); err != nil {
		t.Errorf("expected good channel, err=%v", err)
	}
	c3.Close()

	d.check("4", p, 2, 2)
	p.Close()
}

func TestPoolMonitorCleanup(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}
	c := p.Get()
	c.Send("MONITOR")
	c.Close()

	d.check("", p, 1, 0)
	p.Close()
}

func TestPoolPubSubCleanup(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}

	c := p.Get()
	c.Send("SUBSCRIBE", "x")
	c.Close()

	want := []string{"SUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", "ECHO"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Send("PSUBSCRIBE", "x*")
	c.Close()

	want = []string{"PSUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", "ECHO"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	p.Close()
}

func TestPoolTransactionCleanup(t *testing.T) {
	d := poolDialer{t: t}
	p := &Pool{
		MaxIdle:   2,
		MaxActive: 2,
		Dial:      d.dial,
	}

	c := p.Get()
	c.Do("WATCH", "key")
	c.Do("PING")
	c.Close()

	want := []string{"WATCH", "PING", "UNWATCH"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("UNWATCH")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "UNWATCH", "PING"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("MULTI")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "MULTI", "PING", "DISCARD"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("MULTI")
	c.Do("DISCARD")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "MULTI", "DISCARD", "PING"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	c = p.Get()
	c.Do("WATCH", "key")
	c.Do("MULTI")
	c.Do("EXEC")
	c.Do("PING")
	c.Close()

	want = []string{"WATCH", "MULTI", "EXEC", "PING"}
	if !reflect.DeepEqual(d.commands, want) {
		t.Errorf("got commands %v, want %v", d.commands, want)
	}
	d.commands = nil

	p.Close()
}

func BenchmarkPoolGet(b *testing.B) {
	b.StopTimer()
	p := Pool{Dial: DialTestDB, MaxIdle: 2}
	c := p.Get()
	if err := c.Err(); err != nil {
		b.Fatal(err)
	}
	c.Close()
	defer p.Close()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		c = p.Get()
		c.Close()
	}
}

func BenchmarkPoolGetErr(b *testing.B) {
	b.StopTimer()
	p := Pool{Dial: DialTestDB, MaxIdle: 2}
	c := p.Get()
	if err := c.Err(); err != nil {
		b.Fatal(err)
	}
	c.Close()
	defer p.Close()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		c = p.Get()
		if err := c.Err(); err != nil {
			b.Fatal(err)
		}
		c.Close()
	}
}

func BenchmarkPoolGetPing(b *testing.B) {
	b.StopTimer()
	p := Pool{Dial: DialTestDB, MaxIdle: 2}
	c := p.Get()
	if err := c.Err(); err != nil {
		b.Fatal(err)
	}
	c.Close()
	defer p.Close()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		c = p.Get()
		if _, err := c.Do("PING"); err != nil {
			b.Fatal(err)
		}
		c.Close()
	}
}

const numConcurrent = 10

func BenchmarkPipelineConcurrency(b *testing.B) {
	b.StopTimer()
	c, err := DialTestDB()
	if err != nil {
		b.Fatalf("error connection to database, %v", err)
	}
	defer c.Close()

	var wg sync.WaitGroup
	wg.Add(numConcurrent)

	var pipeline textproto.Pipeline

	b.StartTimer()

	for i := 0; i < numConcurrent; i++ {
		go func() {
			defer wg.Done()
			for i := 0; i < b.N; i++ {
				id := pipeline.Next()
				pipeline.StartRequest(id)
				c.Send("PING")
				c.Flush()
				pipeline.EndRequest(id)
				pipeline.StartResponse(id)
				_, err := c.Receive()
				if err != nil {
					b.Fatal(err)
				}
				c.Flush()
				pipeline.EndResponse(id)
			}
		}()
	}
	wg.Wait()
}

func BenchmarkPoolConcurrency(b *testing.B) {
	b.StopTimer()

	p := Pool{Dial: DialTestDB, MaxIdle: numConcurrent}
	defer p.Close()

	// fill the pool
	conns := make([]Conn, numConcurrent)
	for i := range conns {
		c := p.Get()
		if err := c.Err(); err != nil {
			b.Fatal(err)
		}
		conns[i] = c
	}
	for _, c := range conns {
		c.Close()
	}

	var wg sync.WaitGroup
	wg.Add(numConcurrent)

	b.StartTimer()

	for i := 0; i < numConcurrent; i++ {
		go func() {
			defer wg.Done()
			for i := 0; i < b.N; i++ {
				c := p.Get()
				c.Do("PING")
				c.Close()
			}
		}()
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>/*
       master             board 2 (or ui)
   sub  pub  reply        sub  pub  reply
    |    |                 |    |
    |    -------------------    |
    -----------------------------
*/

package gogadgets

import (
	"encoding/json"
	"fmt"
	"github.com/vaughan0/go-zmq"
	"log"
)

//Sockets fufills the GoGadget interface and is
//added to each Gadget system by App.  It provides
//a way to connect multiple gadgets systems together
//as a single system, and also provides a way for
//an external UI to interact the system.
type Sockets struct {
	host    string
	pubPort int
	subPort int
	ctx     *zmq.Context
	sub     *zmq.Socket
	subChan *zmq.Channels
	pub     *zmq.Socket
	pubChan *zmq.Channels
}

func NewSockets() (*Sockets, error) {
	s := &Sockets{
		host:    "localhost",
		subPort: 6111,
		pubPort: 6112,
	}
	err := s.getMasterSockets()
	return s, err
}

func NewClientSockets(host string) (*Sockets, error) {
	s := &Sockets{
		host:    host,
		subPort: 6111,
		pubPort: 6112,
	}
	err := s.getClientSockets()
	return s, err
}

func (s *Sockets) Send(cmd string) {
	msg := Message{
		Type: COMMAND,
		Body: cmd,
	}
	s.SendMessage(msg)
}

func (s *Sockets) SendMessage(msg Message) {
	b, err := json.Marshal(msg)
	if err != nil {
		fmt.Println("zmq sockets had a problem", err)
	} else {
		s.pub.Send([][]byte{
			[]byte(msg.Type),
			b,
		})
	}
}

func (s *Sockets) Recv() *Message {
	data, err := s.sub.Recv()
	if err != nil {
		panic(err)
	}
	msg := &Message{}
	json.Unmarshal(data[1], msg)
	return msg
}

//Sockets listens for chann Messages from inside the system and
//sends it to external listeners (like a UI), and listens for
//external messages and sends them along to the internal system.
func (s *Sockets) Start(in <-chan Message, out chan<- Message) {
	err := s.getSockets()
	defer s.Close()
	if err != nil {
		log.Println("zmq sockets had a problem", err)
	}
	for {
		select {
		case data := <-s.subChan.In():
			s.sendMessageIn(data, out)
		case msg := <-in:
			s.sendMessageOut(msg)
		case err = <-s.subChan.Errors():
			log.Println(err)
		}
	}
}

//A message that came from inside this gogadgets system
//is sent to outside clients (ui, connected gogadget systems)
func (s *Sockets) sendMessageOut(msg Message) bool {
	keepGoing := true
	if msg.Type == COMMAND && msg.Body == "shutdown" {
		keepGoing = false
	}
	b, err := json.Marshal(msg)
	if err != nil {
		log.Println("zmq sockets had a problem", err)
	} else {
		s.pubChan.Out() <- [][]byte{
			[]byte(msg.Type),
			b,
		}
	}
	return keepGoing
}

//A message that came from outside clients (ui, connected
//gogadget systems) is passed along to this gogadget
//system
func (s *Sockets) sendMessageIn(data [][]byte, out chan<- Message) {
	if len(data) == 2 {
		msg := &Message{}
		json.Unmarshal(data[1], msg)
		msg.Sender = "zmq sockets"
		out <- *msg
	} else {
		log.Println("zmq received an improper message", data)
	}
}

func (s *Sockets) Close() {
	s.sub.Close()
	s.pub.Close()
	if s.subChan != nil {
		s.subChan.Close()
	}
	if s.pubChan != nil {
		s.pubChan.Close()
	}
	s.ctx.Close()
}

func (s *Sockets) getSockets() (err error) {
	if s.host == "localhost" || s.host == "" {
		err = s.getMasterChannels()
	} else {
		err = s.getClientSockets()
		s.subChan = s.sub.Channels()
		s.pubChan = s.pub.Channels()
	}
	return err
}

//Two GoGadgets systems can be joined together into a single system.
//The GoGadgets system that has it's App.Host set as "localhost" is
//the master.  Other systems that wish to join need to be configured
//with the IP address of the master system as App.Host.
func (s *Sockets) getMasterChannels() (err error) {
	err = s.getMasterSockets()
	if err != nil {
		return err
	}
	s.pubChan = s.pub.Channels()
	s.subChan = s.sub.Channels()
	return err
}

func (s *Sockets) getMasterSockets() (err error) {
	s.ctx, err = zmq.NewContext()
	if err != nil {
		return err
	}
	s.pub, err = s.ctx.Socket(zmq.Pub)
	if err != nil {
		return err
	}
	if err = s.pub.Bind(fmt.Sprintf("tcp://*:%d", s.pubPort)); err != nil {
		return err
	}
	sub, err := s.ctx.Socket(zmq.Sub)
	if err != nil {
		return err
	}
	s.sub = sub
	if err = s.sub.Bind(fmt.Sprintf("tcp://*:%d", s.subPort)); err != nil {
		return err
	}
	s.sub.Subscribe([]byte(""))
	return err
}

//This creates the zmq sockets for a GoGadget system that is not the master
//system or a UI.
func (s *Sockets) getClientSockets() (err error) {
	s.ctx, err = zmq.NewContext()
	if err != nil {
		return err
	}
	s.pub, err = s.ctx.Socket(zmq.Pub)
	if err != nil {
		return err
	}
	u := fmt.Sprintf("tcp://%s:%d", s.host, s.subPort)
	if err = s.pub.Connect(u); err != nil {
		return err
	}

	s.sub, err = s.ctx.Socket(zmq.Sub)
	if err != nil {
		return err
	}
	u = fmt.Sprintf("tcp://%s:%d", s.host, s.pubPort)
	if err = s.sub.Connect(u); err != nil {
		return err
	}
	s.sub.Subscribe([]byte(""))
	return err
}

func (s *Sockets) GetUID() string {
	return "zmq sockets"
}
<commit_msg>added status update to zmq sockets<commit_after>/*
       master             board 2 (or ui)
   sub  pub  reply        sub  pub  reply
    |    |                 |    |
    |    -------------------    |
    -----------------------------
*/

package gogadgets

import (
	"encoding/json"
	"fmt"
	"github.com/vaughan0/go-zmq"
	"log"
)

//Sockets fufills the GoGadget interface and is
//added to each Gadget system by App.  It provides
//a way to connect multiple gadgets systems together
//as a single system, and also provides a way for
//an external UI to interact the system.
type Sockets struct {
	master  bool
	host    string
	pubPort int
	subPort int
	ctx     *zmq.Context
	sub     *zmq.Socket
	subChan *zmq.Channels
	pub     *zmq.Socket
	pubChan *zmq.Channels
	updates map[string]Message
}

func NewSockets() (*Sockets, error) {
	s := &Sockets{
		master: true,
		host:    "localhost",
		subPort: 6111,
		pubPort: 6112,
		updates: map[string]Message{},
	}
	err := s.getMasterSockets()
	return s, err
}

func NewClientSockets(host string) (*Sockets, error) {
	s := &Sockets{
		host:    host,
		subPort: 6111,
		pubPort: 6112,
	}
	err := s.getClientSockets()
	return s, err
}

func (s *Sockets) Send(cmd string) {
	msg := Message{
		Type: COMMAND,
		Body: cmd,
	}
	s.SendMessage(msg)
}

func (s *Sockets) SendMessage(msg Message) {
	b, err := json.Marshal(msg)
	if err != nil {
		fmt.Println("zmq sockets had a problem", err)
	} else {
		s.pub.Send([][]byte{
			[]byte(msg.Type),
			b,
		})
	}
}

func (s *Sockets) Recv() *Message {
	data, err := s.sub.Recv()
	if err != nil {
		panic(err)
	}
	msg := &Message{}
	json.Unmarshal(data[1], msg)
	return msg
}

//Sockets listens for chann Messages from inside the system and
//sends it to external listeners (like a UI), and listens for
//external messages and sends them along to the internal system.
func (s *Sockets) Start(in <-chan Message, out chan<- Message) {
	err := s.getSockets()
	defer s.Close()
	if err != nil {
		log.Println("zmq sockets had a problem", err)
	}
	for {
		select {
		case data := <-s.subChan.In():
			s.sendMessageIn(data, out)
		case msg := <-in:
			s.sendMessageOut(msg)
		case err = <-s.subChan.Errors():
			log.Println(err)
		}
	}
}

//A message that came from inside this gogadgets system
//is sent to outside clients (ui, connected gogadget systems)
func (s *Sockets) sendMessageOut(msg Message) bool {
	if s.master && msg.Type == UPDATE {
		s.updates[msg.Sender] = msg
	}
	keepGoing := true
	if msg.Type == COMMAND && msg.Body == "shutdown" {
		keepGoing = false
	}
	b, err := json.Marshal(msg)
	if err != nil {
		log.Println("zmq sockets had a problem", err)
	} else {
		s.pubChan.Out() <- [][]byte{
			[]byte(msg.Type),
			b,
		}
	}
	return keepGoing
}

//A message that came from outside clients (ui, connected
//gogadget systems) is passed along to this gogadget
//system
func (s *Sockets) sendMessageIn(data [][]byte, out chan<- Message) {
	if len(data) == 2 {
		msg := &Message{}
		json.Unmarshal(data[1], msg)
		if msg.Sender == "" {
			msg.Sender = "zmq sockets"
		}
		if s.master && msg.Type == UPDATE {
			s.updates[msg.Sender] = *msg
		}
		if msg.Body == "status" {
			s.sendStatus()
		} else {
			out <- *msg
		}
	} else {
		log.Println("zmq received an improper message", data)
	}
}

//An outside client (like a UI) wants the latest status of
//all gadgets in the system.
func (s *Sockets) sendStatus() {
	b, _ := json.Marshal(s.updates)
	s.pubChan.Out() <- [][]byte{
		[]byte("status"),
		b,
	}
}

func (s *Sockets) Close() {
	s.sub.Close()
	s.pub.Close()
	if s.subChan != nil {
		s.subChan.Close()
	}
	if s.pubChan != nil {
		s.pubChan.Close()
	}
	s.ctx.Close()
}

func (s *Sockets) getSockets() (err error) {
	if s.host == "localhost" || s.host == "" {
		err = s.getMasterChannels()
	} else {
		err = s.getClientSockets()
		s.subChan = s.sub.Channels()
		s.pubChan = s.pub.Channels()
	}
	return err
}

//Two GoGadgets systems can be joined together into a single system.
//The GoGadgets system that has it's App.Host set as "localhost" is
//the master.  Other systems that wish to join need to be configured
//with the IP address of the master system as App.Host.
func (s *Sockets) getMasterChannels() (err error) {
	err = s.getMasterSockets()
	if err != nil {
		return err
	}
	s.pubChan = s.pub.Channels()
	s.subChan = s.sub.Channels()
	return err
}

func (s *Sockets) getMasterSockets() (err error) {
	s.ctx, err = zmq.NewContext()
	if err != nil {
		return err
	}
	s.pub, err = s.ctx.Socket(zmq.Pub)
	if err != nil {
		return err
	}
	if err = s.pub.Bind(fmt.Sprintf("tcp://*:%d", s.pubPort)); err != nil {
		return err
	}
	sub, err := s.ctx.Socket(zmq.Sub)
	if err != nil {
		return err
	}
	s.sub = sub
	if err = s.sub.Bind(fmt.Sprintf("tcp://*:%d", s.subPort)); err != nil {
		return err
	}
	s.sub.Subscribe([]byte(""))
	return err
}

//This creates the zmq sockets for a GoGadget system that is not the master
//system or a UI.
func (s *Sockets) getClientSockets() (err error) {
	s.ctx, err = zmq.NewContext()
	if err != nil {
		return err
	}
	s.pub, err = s.ctx.Socket(zmq.Pub)
	if err != nil {
		return err
	}
	u := fmt.Sprintf("tcp://%s:%d", s.host, s.subPort)
	if err = s.pub.Connect(u); err != nil {
		return err
	}

	s.sub, err = s.ctx.Socket(zmq.Sub)
	if err != nil {
		return err
	}
	u = fmt.Sprintf("tcp://%s:%d", s.host, s.pubPort)
	if err = s.sub.Connect(u); err != nil {
		return err
	}
	s.sub.Subscribe([]byte(""))
	return err
}

func (s *Sockets) GetUID() string {
	return "zmq sockets"
}
<|endoftext|>


<commit_before>package storage

import (
	"crypto/rand"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/coreos/etcd/storage/storagepb"
)

// TODO: improve to a unit test
func TestRangeLimitWhenKeyDeleted(t *testing.T) {
	s := newStore(tmpPath)
	defer os.Remove(tmpPath)

	s.Put([]byte("foo"), []byte("bar"))
	s.Put([]byte("foo1"), []byte("bar1"))
	s.Put([]byte("foo2"), []byte("bar2"))
	s.DeleteRange([]byte("foo1"), nil)
	kvs := []storagepb.KeyValue{
		{Key: []byte("foo"), Value: []byte("bar"), CreateIndex: 1, ModIndex: 1, Version: 1},
		{Key: []byte("foo2"), Value: []byte("bar2"), CreateIndex: 3, ModIndex: 3, Version: 1},
	}

	tests := []struct {
		limit int64
		wkvs  []storagepb.KeyValue
	}{
		// no limit
		{0, kvs},
		{1, kvs[:1]},
		{2, kvs},
		{3, kvs},
	}
	for i, tt := range tests {
		kvs, _, err := s.Range([]byte("foo"), []byte("foo3"), tt.limit, 0)
		if err != nil {
			t.Fatalf("#%d: range error (%v)", i, err)
		}
		if !reflect.DeepEqual(kvs, tt.wkvs) {
			t.Errorf("#%d: kvs = %+v, want %+v", i, kvs, tt.wkvs)
		}
	}
}

func TestRestoreContinueUnfinishedCompaction(t *testing.T) {
	s0 := newStore(tmpPath)
	defer os.Remove(tmpPath)

	s0.Put([]byte("foo"), []byte("bar"))
	s0.Put([]byte("foo"), []byte("bar1"))
	s0.Put([]byte("foo"), []byte("bar2"))

	// write scheduled compaction, but not do compaction
	rbytes := newRevBytes()
	revToBytes(revision{main: 2}, rbytes)
	tx := s0.b.BatchTx()
	tx.Lock()
	tx.UnsafePut(metaBucketName, scheduledCompactKeyName, rbytes)
	tx.Unlock()

	s0.Close()

	s1 := newStore(tmpPath)
	s1.Restore()

	// wait for scheduled compaction to be finished
	time.Sleep(100 * time.Millisecond)

	if _, _, err := s1.Range([]byte("foo"), nil, 0, 2); err != ErrCompacted {
		t.Errorf("range on compacted rev error = %v, want %v", err, ErrCompacted)
	}
	// check the key in backend is deleted
	revbytes := newRevBytes()
	// TODO: compact should delete main=2 key too
	revToBytes(revision{main: 1}, revbytes)
	tx = s1.b.BatchTx()
	tx.Lock()
	ks, _ := tx.UnsafeRange(keyBucketName, revbytes, nil, 0)
	if len(ks) != 0 {
		t.Errorf("key for rev %+v still exists, want deleted", bytesToRev(revbytes))
	}
	tx.Unlock()
}

func BenchmarkStorePut(b *testing.B) {
	s := newStore(tmpPath)
	defer os.Remove(tmpPath)

	// prepare keys
	keys := make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		keys[i] = make([]byte, 64)
		rand.Read(keys[i])
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Put(keys[i], []byte("foo"))
	}
}
<commit_msg>storage: add mock tests for store struct<commit_after>package storage

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"io"
	"math"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/coreos/etcd/pkg/testutil"
	"github.com/coreos/etcd/storage/backend"
	"github.com/coreos/etcd/storage/storagepb"
)

func TestStorePut(t *testing.T) {
	tests := []struct {
		rev revision
		r   indexGetResp

		wrev    revision
		wev     storagepb.Event
		wputrev revision
	}{
		{
			revision{1, 0},
			indexGetResp{revision{}, revision{}, 0, ErrRevisionNotFound},
			revision{1, 1},
			storagepb.Event{
				Type: storagepb.PUT,
				Kv: &storagepb.KeyValue{
					Key:         []byte("foo"),
					Value:       []byte("bar"),
					CreateIndex: 2,
					ModIndex:    2,
					Version:     1,
				},
			},
			revision{2, 0},
		},
		{
			revision{1, 1},
			indexGetResp{revision{2, 0}, revision{2, 0}, 1, nil},
			revision{1, 2},
			storagepb.Event{
				Type: storagepb.PUT,
				Kv: &storagepb.KeyValue{
					Key:         []byte("foo"),
					Value:       []byte("bar"),
					CreateIndex: 2,
					ModIndex:    2,
					Version:     2,
				},
			},
			revision{2, 1},
		},
		{
			revision{2, 0},
			indexGetResp{revision{2, 1}, revision{2, 0}, 2, nil},
			revision{2, 1},
			storagepb.Event{
				Type: storagepb.PUT,
				Kv: &storagepb.KeyValue{
					Key:         []byte("foo"),
					Value:       []byte("bar"),
					CreateIndex: 2,
					ModIndex:    3,
					Version:     3,
				},
			},
			revision{3, 0},
		},
	}
	for i, tt := range tests {
		s, b, index := newFakeStore()
		s.currentRev = tt.rev
		index.indexGetRespc <- tt.r

		s.put([]byte("foo"), []byte("bar"))

		data, err := tt.wev.Marshal()
		if err != nil {
			t.Errorf("#%d: marshal err = %v, want nil", i, err)
		}
		wact := []testutil.Action{
			{"put", []interface{}{keyBucketName, newTestBytes(tt.wputrev), data}},
		}
		if g := b.tx.Action(); !reflect.DeepEqual(g, wact) {
			t.Errorf("#%d: tx action = %+v, want %+v", i, g, wact)
		}
		wact = []testutil.Action{
			{"get", []interface{}{[]byte("foo"), tt.wputrev.main}},
			{"put", []interface{}{[]byte("foo"), tt.wputrev}},
		}
		if g := index.Action(); !reflect.DeepEqual(g, wact) {
			t.Errorf("#%d: index action = %+v, want %+v", i, g, wact)
		}
		if s.currentRev != tt.wrev {
			t.Errorf("#%d: rev = %+v, want %+v", i, s.currentRev, tt.wrev)
		}
	}
}

func TestStoreRange(t *testing.T) {
	ev := storagepb.Event{
		Type: storagepb.PUT,
		Kv: &storagepb.KeyValue{
			Key:         []byte("foo"),
			Value:       []byte("bar"),
			CreateIndex: 1,
			ModIndex:    2,
			Version:     1,
		},
	}
	evb, err := ev.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	currev := revision{1, 1}
	wrev := int64(2)

	tests := []struct {
		idxr indexRangeResp
		r    rangeResp
	}{
		{
			indexRangeResp{[][]byte{[]byte("foo")}, []revision{{2, 0}}},
			rangeResp{[][]byte{newTestBytes(revision{2, 0})}, [][]byte{evb}},
		},
		{
			indexRangeResp{[][]byte{[]byte("foo"), []byte("foo1")}, []revision{{2, 0}, {3, 0}}},
			rangeResp{[][]byte{newTestBytes(revision{2, 0})}, [][]byte{evb}},
		},
	}
	for i, tt := range tests {
		s, b, index := newFakeStore()
		s.currentRev = currev
		b.tx.rangeRespc <- tt.r
		index.indexRangeRespc <- tt.idxr

		kvs, rev, err := s.rangeKeys([]byte("foo"), []byte("goo"), 1, 0)
		if err != nil {
			t.Errorf("#%d: err = %v, want nil", i, err)
		}
		if w := []storagepb.KeyValue{*ev.Kv}; !reflect.DeepEqual(kvs, w) {
			t.Errorf("#%d: kvs = %+v, want %+v", i, kvs, w)
		}
		if rev != wrev {
			t.Errorf("#%d: rev = %d, want %d", i, rev, wrev)
		}

		wact := []testutil.Action{
			{"range", []interface{}{keyBucketName, newTestBytes(tt.idxr.revs[0]), []byte(nil), int64(0)}},
		}
		if g := b.tx.Action(); !reflect.DeepEqual(g, wact) {
			t.Errorf("#%d: tx action = %+v, want %+v", i, g, wact)
		}
		wact = []testutil.Action{
			{"range", []interface{}{[]byte("foo"), []byte("goo"), wrev}},
		}
		if g := index.Action(); !reflect.DeepEqual(g, wact) {
			t.Errorf("#%d: index action = %+v, want %+v", i, g, wact)
		}
		if s.currentRev != currev {
			t.Errorf("#%d: current rev = %+v, want %+v", i, s.currentRev, currev)
		}
	}
}

func TestStoreDeleteRange(t *testing.T) {
	tests := []struct {
		rev revision
		r   indexRangeResp

		wrev    revision
		wrrev   int64
		wdelrev revision
	}{
		{
			revision{2, 0},
			indexRangeResp{[][]byte{[]byte("foo")}, []revision{{2, 0}}},
			revision{2, 1},
			2,
			revision{3, 0},
		},
		{
			revision{2, 1},
			indexRangeResp{[][]byte{[]byte("foo")}, []revision{{2, 0}}},
			revision{2, 2},
			3,
			revision{3, 1},
		},
	}
	for i, tt := range tests {
		s, b, index := newFakeStore()
		s.currentRev = tt.rev
		index.indexRangeRespc <- tt.r

		n := s.deleteRange([]byte("foo"), []byte("goo"))
		if n != 1 {
			t.Errorf("#%d: n = %d, want 1", i, n)
		}

		data, err := (&storagepb.Event{
			Type: storagepb.DELETE,
			Kv: &storagepb.KeyValue{
				Key: []byte("foo"),
			},
		}).Marshal()
		if err != nil {
			t.Errorf("#%d: marshal err = %v, want nil", i, err)
		}
		wact := []testutil.Action{
			{"put", []interface{}{keyBucketName, newTestBytes(tt.wdelrev), data}},
		}
		if g := b.tx.Action(); !reflect.DeepEqual(g, wact) {
			t.Errorf("#%d: tx action = %+v, want %+v", i, g, wact)
		}
		wact = []testutil.Action{
			{"range", []interface{}{[]byte("foo"), []byte("goo"), tt.wrrev}},
			{"tombstone", []interface{}{[]byte("foo"), tt.wdelrev}},
		}
		if g := index.Action(); !reflect.DeepEqual(g, wact) {
			t.Errorf("#%d: index action = %+v, want %+v", i, g, wact)
		}
		if s.currentRev != tt.wrev {
			t.Errorf("#%d: rev = %+v, want %+v", i, s.currentRev, tt.wrev)
		}
	}
}

func TestStoreCompact(t *testing.T) {
	s, b, index := newFakeStore()
	s.currentRev = revision{3, 0}
	index.indexCompactRespc <- map[revision]struct{}{revision{1, 0}: {}}
	b.tx.rangeRespc <- rangeResp{[][]byte{newTestBytes(revision{1, 0}), newTestBytes(revision{2, 0})}, nil}

	s.Compact(3)
	s.wg.Wait()

	if s.compactMainRev != 3 {
		t.Errorf("compact main rev = %d, want 3", s.compactMainRev)
	}
	end := make([]byte, 8)
	binary.BigEndian.PutUint64(end, uint64(4))
	wact := []testutil.Action{
		{"put", []interface{}{metaBucketName, scheduledCompactKeyName, newTestBytes(revision{3, 0})}},
		{"range", []interface{}{keyBucketName, make([]byte, 17), end, int64(10000)}},
		{"delete", []interface{}{keyBucketName, newTestBytes(revision{2, 0})}},
		{"put", []interface{}{metaBucketName, finishedCompactKeyName, newTestBytes(revision{3, 0})}},
	}
	if g := b.tx.Action(); !reflect.DeepEqual(g, wact) {
		t.Errorf("tx actions = %+v, want %+v", g, wact)
	}
	wact = []testutil.Action{
		{"compact", []interface{}{int64(3)}},
	}
	if g := index.Action(); !reflect.DeepEqual(g, wact) {
		t.Errorf("index action = %+v, want %+v", g, wact)
	}
}

func TestStoreRestore(t *testing.T) {
	s, b, index := newFakeStore()

	putev := storagepb.Event{
		Type: storagepb.PUT,
		Kv: &storagepb.KeyValue{
			Key:         []byte("foo"),
			Value:       []byte("bar"),
			CreateIndex: 3,
			ModIndex:    3,
			Version:     1,
		},
	}
	putevb, err := putev.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	delev := storagepb.Event{
		Type: storagepb.DELETE,
		Kv: &storagepb.KeyValue{
			Key: []byte("foo"),
		},
	}
	delevb, err := delev.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	b.tx.rangeRespc <- rangeResp{[][]byte{finishedCompactKeyName}, [][]byte{newTestBytes(revision{2, 0})}}
	b.tx.rangeRespc <- rangeResp{[][]byte{newTestBytes(revision{3, 0}), newTestBytes(revision{4, 0})}, [][]byte{putevb, delevb}}
	b.tx.rangeRespc <- rangeResp{[][]byte{scheduledCompactKeyName}, [][]byte{newTestBytes(revision{2, 0})}}

	s.Restore()

	if s.compactMainRev != 2 {
		t.Errorf("compact rev = %d, want 4", s.compactMainRev)
	}
	wrev := revision{4, 0}
	if !reflect.DeepEqual(s.currentRev, wrev) {
		t.Errorf("current rev = %v, want %v", s.currentRev, wrev)
	}
	wact := []testutil.Action{
		{"range", []interface{}{metaBucketName, finishedCompactKeyName, []byte(nil), int64(0)}},
		{"range", []interface{}{keyBucketName, newTestBytes(revision{}), newTestBytes(revision{math.MaxInt64, math.MaxInt64}), int64(0)}},
		{"range", []interface{}{metaBucketName, scheduledCompactKeyName, []byte(nil), int64(0)}},
	}
	if g := b.tx.Action(); !reflect.DeepEqual(g, wact) {
		t.Errorf("tx actions = %+v, want %+v", g, wact)
	}
	wact = []testutil.Action{
		{"restore", []interface{}{[]byte("foo"), revision{3, 0}, revision{3, 0}, int64(1)}},
		{"tombstone", []interface{}{[]byte("foo"), revision{4, 0}}},
	}
	if g := index.Action(); !reflect.DeepEqual(g, wact) {
		t.Errorf("index action = %+v, want %+v", g, wact)
	}
}

func TestRestoreContinueUnfinishedCompaction(t *testing.T) {
	s0 := newStore(tmpPath)
	defer os.Remove(tmpPath)

	s0.Put([]byte("foo"), []byte("bar"))
	s0.Put([]byte("foo"), []byte("bar1"))
	s0.Put([]byte("foo"), []byte("bar2"))

	// write scheduled compaction, but not do compaction
	rbytes := newRevBytes()
	revToBytes(revision{main: 2}, rbytes)
	tx := s0.b.BatchTx()
	tx.Lock()
	tx.UnsafePut(metaBucketName, scheduledCompactKeyName, rbytes)
	tx.Unlock()

	s0.Close()

	s1 := newStore(tmpPath)
	s1.Restore()

	// wait for scheduled compaction to be finished
	time.Sleep(100 * time.Millisecond)

	if _, _, err := s1.Range([]byte("foo"), nil, 0, 2); err != ErrCompacted {
		t.Errorf("range on compacted rev error = %v, want %v", err, ErrCompacted)
	}
	// check the key in backend is deleted
	revbytes := newRevBytes()
	// TODO: compact should delete main=2 key too
	revToBytes(revision{main: 1}, revbytes)
	tx = s1.b.BatchTx()
	tx.Lock()
	ks, _ := tx.UnsafeRange(keyBucketName, revbytes, nil, 0)
	if len(ks) != 0 {
		t.Errorf("key for rev %+v still exists, want deleted", bytesToRev(revbytes))
	}
	tx.Unlock()
}

func BenchmarkStorePut(b *testing.B) {
	s := newStore(tmpPath)
	defer os.Remove(tmpPath)

	// prepare keys
	keys := make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		keys[i] = make([]byte, 64)
		rand.Read(keys[i])
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Put(keys[i], []byte("foo"))
	}
}

func newTestBytes(rev revision) []byte {
	bytes := newRevBytes()
	revToBytes(rev, bytes)
	return bytes
}

func newFakeStore() (*store, *fakeBackend, *fakeIndex) {
	b := &fakeBackend{&fakeBatchTx{rangeRespc: make(chan rangeResp, 5)}}
	index := &fakeIndex{
		indexGetRespc:     make(chan indexGetResp, 1),
		indexRangeRespc:   make(chan indexRangeResp, 1),
		indexCompactRespc: make(chan map[revision]struct{}, 1),
	}
	return &store{
		b:              b,
		kvindex:        index,
		currentRev:     revision{},
		compactMainRev: -1,
	}, b, index
}

type rangeResp struct {
	keys [][]byte
	vals [][]byte
}

type fakeBatchTx struct {
	testutil.Recorder
	rangeRespc chan rangeResp
}

func (b *fakeBatchTx) Lock()                          {}
func (b *fakeBatchTx) Unlock()                        {}
func (b *fakeBatchTx) UnsafeCreateBucket(name []byte) {}
func (b *fakeBatchTx) UnsafePut(bucketName []byte, key []byte, value []byte) {
	b.Recorder.Record(testutil.Action{Name: "put", Params: []interface{}{bucketName, key, value}})
}
func (b *fakeBatchTx) UnsafeRange(bucketName []byte, key, endKey []byte, limit int64) (keys [][]byte, vals [][]byte) {
	b.Recorder.Record(testutil.Action{Name: "range", Params: []interface{}{bucketName, key, endKey, limit}})
	r := <-b.rangeRespc
	return r.keys, r.vals
}
func (b *fakeBatchTx) UnsafeDelete(bucketName []byte, key []byte) {
	b.Recorder.Record(testutil.Action{Name: "delete", Params: []interface{}{bucketName, key}})
}
func (b *fakeBatchTx) Commit()        {}
func (b *fakeBatchTx) CommitAndStop() {}

type fakeBackend struct {
	tx *fakeBatchTx
}

func (b *fakeBackend) BatchTx() backend.BatchTx                  { return b.tx }
func (b *fakeBackend) Snapshot(w io.Writer) (n int64, err error) { return 0, errors.New("unsupported") }
func (b *fakeBackend) ForceCommit()                              {}
func (b *fakeBackend) Close() error                              { return nil }

type indexGetResp struct {
	rev     revision
	created revision
	ver     int64
	err     error
}

type indexRangeResp struct {
	keys [][]byte
	revs []revision
}

type fakeIndex struct {
	testutil.Recorder
	indexGetRespc     chan indexGetResp
	indexRangeRespc   chan indexRangeResp
	indexCompactRespc chan map[revision]struct{}
}

func (i *fakeIndex) Get(key []byte, atRev int64) (rev, created revision, ver int64, err error) {
	i.Recorder.Record(testutil.Action{Name: "get", Params: []interface{}{key, atRev}})
	r := <-i.indexGetRespc
	return r.rev, r.created, r.ver, r.err
}
func (i *fakeIndex) Range(key, end []byte, atRev int64) ([][]byte, []revision) {
	i.Recorder.Record(testutil.Action{Name: "range", Params: []interface{}{key, end, atRev}})
	r := <-i.indexRangeRespc
	return r.keys, r.revs
}
func (i *fakeIndex) Put(key []byte, rev revision) {
	i.Recorder.Record(testutil.Action{Name: "put", Params: []interface{}{key, rev}})
}
func (i *fakeIndex) Restore(key []byte, created, modified revision, ver int64) {
	i.Recorder.Record(testutil.Action{Name: "restore", Params: []interface{}{key, created, modified, ver}})
}
func (i *fakeIndex) Tombstone(key []byte, rev revision) error {
	i.Recorder.Record(testutil.Action{Name: "tombstone", Params: []interface{}{key, rev}})
	return nil
}
func (i *fakeIndex) Compact(rev int64) map[revision]struct{} {
	i.Recorder.Record(testutil.Action{Name: "compact", Params: []interface{}{rev}})
	return <-i.indexCompactRespc
}
func (i *fakeIndex) Equal(b index) bool { return false }
<|endoftext|>


<commit_before>package request

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"log/syslog"
	"math/rand"
	"net/http"
	"os"
	"reflect"
	"strings"
	"time"

	"github.com/STNS/STNS/attribute"
	"github.com/STNS/libnss_stns/config"
)

var ConfigFileName = "/etc/stns/libnss_stns.conf"
var Loaded *config.Config

type Request struct {
	apiPath string
	Config  *config.Config
}

func choice(s []string) string {
	rand.Seed(time.Now().UnixNano())
	i := rand.Intn(len(s))
	return s[i]
}

func NewRequest(resource string, column string, value string) (*Request, error) {
	r := Request{}
	if err := r.Init(); err != nil {
		return nil, err
	}
	urls := []string{resource, column}

	if value != "" {
		urls = append(urls, value)
	}

	r.apiPath = strings.Join(urls, "/")
	return &r, nil
}

func (r *Request) Get() (attribute.UserGroups, error) {
	var resultErr error
	client := &http.Client{}

	for _, endPoint := range r.Config.ApiEndPoint {
		req, err := http.NewRequest("GET", endPoint+"/"+r.apiPath, nil)
		if err != nil {
			resultErr = err
			continue
		}

		res, err := client.Do(req)
		if err != nil {
			resultErr = err
			continue
		}

		defer res.Body.Close()

		if res.StatusCode != http.StatusOK {
			continue
		}

		body, err := ioutil.ReadAll(res.Body)
		if err != nil {
			resultErr = err
			continue
		}

		var attr attribute.UserGroups
		err = json.Unmarshal(body, &attr)
		if err != nil {
			resultErr = err
			continue
		}
		return attr, nil
	}
	return nil, resultErr
}

func (r *Request) Init() error {

	if reflect.ValueOf(Loaded).IsNil() {
		logger, err := syslog.New(syslog.LOG_ERR|syslog.LOG_USER, os.Args[0])
		if err != nil {
			// syslog not found
			fmt.Print(err)
		} else {
			log.SetOutput(logger)
		}
		if ConfigFileName != "" {
			config, err := config.Load(ConfigFileName)
			if err != nil {
				log.Print(err)
				return err
			}
			Loaded = config
		} else {
			Loaded = &config.Config{}
		}
	}
	r.Config = Loaded
	return nil
}
<commit_msg>assync request<commit_after>package request

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"log/syslog"
	"net/http"
	"os"
	"reflect"
	"strings"

	"github.com/STNS/STNS/attribute"
	"github.com/STNS/libnss_stns/config"
)

var ConfigFileName = "/etc/stns/libnss_stns.conf"
var Loaded *config.Config

type Request struct {
	apiPath string
	Config  *config.Config
}

type HttpResponse struct {
	response *http.Response
	err      error
}

func NewRequest(resource string, column string, value string) (*Request, error) {
	r := Request{}
	if err := r.Init(); err != nil {
		return nil, err
	}
	urls := []string{resource, column}

	if value != "" {
		urls = append(urls, value)
	}

	r.apiPath = strings.Join(urls, "/")
	return &r, nil
}

func (r *Request) asyncHttpGets() ([]byte, error) {
	responses := []*HttpResponse{}
	ch := make(chan *HttpResponse, len(r.Config.ApiEndPoint))
	for _, endPoint := range r.Config.ApiEndPoint {
		go func(endPoint string) {
			url := endPoint + "/" + r.apiPath
			res, err := http.Get(url)

			if err != nil {
				ch <- &HttpResponse{nil, err}
			} else {
				ch <- &HttpResponse{res, err}
			}
		}(endPoint)
	}

	for {
		select {
		case c := <-ch:
			responses = append(responses, c)
			if c.response != nil {
				if c.response.StatusCode != http.StatusOK {
					continue
				}
				body, err := ioutil.ReadAll(c.response.Body)
				if err != nil {
					continue
				}
				c.response.Body.Close()
				return body, nil
			} else if c.response == nil && len(responses) == len(r.Config.ApiEndPoint) {
				return nil, c.err
			}
		}
	}
}
func (r *Request) Get() (attribute.UserGroups, error) {
	body, err := r.asyncHttpGets()
	if err != nil {
		return nil, err
	}

	var attr attribute.UserGroups
	err = json.Unmarshal(body, &attr)
	if err != nil {
		return nil, err
	}
	return attr, nil
}

func (r *Request) Init() error {

	if reflect.ValueOf(Loaded).IsNil() {
		logger, err := syslog.New(syslog.LOG_ERR|syslog.LOG_USER, os.Args[0])
		if err != nil {
			// syslog not found
			fmt.Print(err)
		} else {
			log.SetOutput(logger)
		}
		if ConfigFileName != "" {
			config, err := config.Load(ConfigFileName)
			if err != nil {
				log.Print(err)
				return err
			}
			Loaded = config
		} else {
			Loaded = &config.Config{}
		}
	}
	r.Config = Loaded
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
)

type Config struct {
	Worker int `json:"worker"`
	Tasks  int `json:"tasks"`
}

func main() {
	var conf Config
	dat, err := ioutil.ReadFile("read-json-config-file.conf")
	if err != nil {
		panic(err)
	}
	fmt.Println("dat", dat)

	if err = json.Unmarshal(dat, &conf); err != nil {
		panic(err)
	}
	fmt.Println("conf", conf)
}
<commit_msg>print Config<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
)

type Config struct {
	Worker int `json:"worker"`
	Tasks  int `json:"tasks"`
}

func main() {
	var conf Config
	dat, err := ioutil.ReadFile("read-json-config-file.conf")
	if err != nil {
		panic(err)
	}
	fmt.Println("dat", string(dat))

	if err = json.Unmarshal(dat, &conf); err != nil {
		panic(err)
	}
	fmt.Println("conf: Worker ", conf.Worker)
	fmt.Println("conf: Tasks  ", conf.Tasks)
}
<|endoftext|>


<commit_before>// Copyright 2015-2016, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 storemock_test

import (
	"testing"

	"github.com/corestoreio/csfw/config/cfgmock"
	"github.com/corestoreio/csfw/config/cfgpath"
	"github.com/corestoreio/csfw/store/scope"
	"github.com/corestoreio/csfw/store/storemock"
	"github.com/stretchr/testify/assert"
)

func TestMustNewStoreAU_ConfigNil(t *testing.T) {
	sAU := storemock.MustNewStoreAU(cfgmock.NewService())
	assert.NotNil(t, sAU)
	assert.NotNil(t, sAU.Config)
	assert.NotNil(t, sAU.Website.Config)

	assert.Exactly(t, int64(5), sAU.Config.StoreID)
	assert.Exactly(t, int64(2), sAU.Config.WebsiteID)

	assert.Exactly(t, int64(0), sAU.Website.Config.StoreID)
	assert.Exactly(t, int64(2), sAU.Website.Config.WebsiteID)

}

func TestMustNewStoreAU_ConfigNonNil(t *testing.T) {
	sAU := storemock.MustNewStoreAU(cfgmock.NewService())
	assert.NotNil(t, sAU)
	assert.NotNil(t, sAU.Config)
	assert.NotNil(t, sAU.Website.Config)
}

func TestMustNewStoreAU_Config(t *testing.T) {
	var configPath = cfgpath.MustNewByParts("aa/bb/cc")

	sm := cfgmock.NewService(cfgmock.PathValue{
		configPath.String():                "DefaultScopeString",
		configPath.BindWebsite(2).String(): "WebsiteScopeString",
		configPath.BindStore(5).String():   "StoreScopeString",
	})
	aust := storemock.MustNewStoreAU(sm)

	haveS, err := aust.Website.Config.String(configPath.Route)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "WebsiteScopeString", haveS)

	haveS, err = aust.Website.Config.String(configPath.Route, scope.Default)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "DefaultScopeString", haveS)

	haveS, err = aust.Config.String(configPath.Route)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "StoreScopeString", haveS)

	haveS, err = aust.Config.String(configPath.Route, scope.Default)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "DefaultScopeString", haveS)

	assert.Exactly(t, scope.TypeIDs{scope.DefaultTypeID, scope.Website.Pack(2), scope.Store.Pack(5)}, sm.AllInvocations().TypeIDs())
	assert.Exactly(t, 3, sm.AllInvocations().PathCount())

}
<commit_msg>store/storemock: Adjust API to package config/cfgmock.<commit_after>// Copyright 2015-2016, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 storemock_test

import (
	"testing"

	"github.com/corestoreio/csfw/config/cfgmock"
	"github.com/corestoreio/csfw/config/cfgpath"
	"github.com/corestoreio/csfw/store/scope"
	"github.com/corestoreio/csfw/store/storemock"
	"github.com/stretchr/testify/assert"
)

func TestMustNewStoreAU_ConfigNil(t *testing.T) {
	sAU := storemock.MustNewStoreAU(cfgmock.NewService())
	assert.NotNil(t, sAU)
	assert.NotNil(t, sAU.Config)
	assert.NotNil(t, sAU.Website.Config)

	assert.Exactly(t, int64(5), sAU.Config.StoreID)
	assert.Exactly(t, int64(2), sAU.Config.WebsiteID)

	assert.Exactly(t, int64(0), sAU.Website.Config.StoreID)
	assert.Exactly(t, int64(2), sAU.Website.Config.WebsiteID)

}

func TestMustNewStoreAU_ConfigNonNil(t *testing.T) {
	sAU := storemock.MustNewStoreAU(cfgmock.NewService())
	assert.NotNil(t, sAU)
	assert.NotNil(t, sAU.Config)
	assert.NotNil(t, sAU.Website.Config)
}

func TestMustNewStoreAU_Config(t *testing.T) {
	var configPath = cfgpath.MustNewByParts("aa/bb/cc")

	sm := cfgmock.NewService(cfgmock.PathValue{
		configPath.String():                "DefaultScopeString",
		configPath.BindWebsite(2).String(): "WebsiteScopeString",
		configPath.BindStore(5).String():   "StoreScopeString",
	})
	aust := storemock.MustNewStoreAU(sm)

	haveS, err := aust.Website.Config.String(configPath.Route)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "WebsiteScopeString", haveS)

	haveS, err = aust.Website.Config.String(configPath.Route, scope.Default)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "DefaultScopeString", haveS)

	haveS, err = aust.Config.String(configPath.Route)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "StoreScopeString", haveS)

	haveS, err = aust.Config.String(configPath.Route, scope.Default)
	if err != nil {
		t.Fatalf("%+v", err)
	}
	assert.Exactly(t, "DefaultScopeString", haveS)

	assert.Exactly(t, scope.TypeIDs{scope.DefaultTypeID, scope.Website.Pack(2), scope.Store.Pack(5)}, sm.AllInvocations().ScopeIDs())
	assert.Exactly(t, 3, sm.AllInvocations().PathCount())

}
<|endoftext|>