<commit_before>// This program scrapes the cloudformation documentation to determine the schema
// and produces a go program to the file specified by the `-out` flag.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"regexp"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"github.com/mitchellh/go-wordwrap"
)

// TemplateReference describes the CloudFormation template schema
type TemplateReference struct {
	Resources []*Resource
}

func (tr *TemplateReference) Load() error {
	for _, docURI := range []string{"aws-template-resource-type-ref.html", "aws-product-property-reference.html"} {
		docReader, err := getDoc(docURI)
		if err != nil {
			return err
		}
		doc, err := goquery.NewDocumentFromReader(docReader)
		if err != nil {
			return err
		}
		doc.Find(".highlights li a").Each(func(i int, s *goquery.Selection) {
			name := s.Text()
			name = regexp.MustCompile("\\s+").ReplaceAllString(name, " ")
			href, _ := s.Attr("href")
			tr.Resources = append(tr.Resources, &Resource{Name: name, Href: href})

		})
	}

	for _, resource := range tr.Resources {
		if err := resource.Load(); err != nil {
			return fmt.Errorf("%s: %s", resource.Name, err)
		}
	}
	return nil
}

func (tr *TemplateReference) WriteGo(w io.Writer) {
	fmt.Fprintf(w, "package cloudformation\n")
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "import \"time\"\n")
	fmt.Fprintf(w, "import \"encoding/json\"\n")
	for _, resource := range tr.Resources {
		fmt.Fprintf(w, "\n")
		fmt.Fprintf(w, "// %s represents %s\n", resource.GoName(), resource.Name)
		fmt.Fprintf(w, "//\n")
		fmt.Fprintf(w, "// see %s%s\n", rootURL, resource.Href)
		fmt.Fprintf(w, "type %s struct {\n", resource.GoName())
		for i, p := range resource.Properties {
			if i > 0 {
				fmt.Fprintf(w, "\n")
			}
			fmt.Fprintf(w, "%s", p.Comment("  // "))
			fmt.Fprintf(w, "  %s %s `json:\"%s,omitempty\"`\n", p.GoName(), p.GoType(tr), p.Name)
		}
		fmt.Fprintf(w, "}\n")
		fmt.Fprintf(w, "\n")
		if resource.IsTopLevelResource() {

			fmt.Fprintf(w, "// ResourceType returns %s to implement the ResourceProperties interface\n", resource.Name)
			fmt.Fprintf(w, "func (s %s) ResourceType() string {\n", resource.GoName())
			fmt.Fprintf(w, "	return %q\n", resource.Name)
			fmt.Fprintf(w, "}\n")
			fmt.Fprintf(w, "\n")
		}

		// Cloudformation allows a single object when a list of objects is expected. To
		// handle this we need to generate a *List type. This applies mostly only to
		// non-top-level objects. (The only exception is AWS::Route53::RecordSet which is
		// both a top level resource and a child element of AWS::Route53::RecordSetGroup
		if !resource.IsTopLevelResource() || resource.GoName() == "Route53RecordSet" {
			fmt.Fprintf(w, "// %sList represents a list of %s\n", resource.GoName(), resource.GoName())
			fmt.Fprintf(w, "type %sList []%s\n", resource.GoName(), resource.GoName())
			fmt.Fprintf(w, "\n")
			fmt.Fprintf(w, "// UnmarshalJSON sets the object from the provided JSON representation\n")
			fmt.Fprintf(w, "func (l *%sList) UnmarshalJSON(buf []byte) error {\n", resource.GoName())
			fmt.Fprintf(w, "	// Cloudformation allows a single object when a list of objects is expected\n")
			fmt.Fprintf(w, "	item := %s{}\n", resource.GoName())
			fmt.Fprintf(w, "	if err := json.Unmarshal(buf, &item); err == nil {\n")
			fmt.Fprintf(w, "		*l = %sList{item}\n", resource.GoName())
			fmt.Fprintf(w, "		return nil\n")
			fmt.Fprintf(w, "	}\n")
			fmt.Fprintf(w, "	list := []%s{}\n", resource.GoName())
			fmt.Fprintf(w, "	err := json.Unmarshal(buf, &list)\n")
			fmt.Fprintf(w, "	if err == nil {\n")
			fmt.Fprintf(w, "		*l = %sList(list)\n", resource.GoName())
			fmt.Fprintf(w, "		return nil\n")
			fmt.Fprintf(w, "	}\n")
			fmt.Fprintf(w, "	return err\n")
			fmt.Fprintf(w, "}\n")
			fmt.Fprintf(w, "\n")
		}
	}

	fmt.Fprintf(w, "// NewResourceByType returns a new resource object correspoding with the provided type\n")
	fmt.Fprintf(w, "func NewResourceByType(typeName string) ResourceProperties {\n")
	fmt.Fprintf(w, "	switch typeName {\n")

	for _, resource := range tr.Resources {
		if !resource.IsTopLevelResource() {
			continue
		}
		fmt.Fprintf(w, "		case %q:\n", resource.Name)
		fmt.Fprintf(w, "			return &%s{}\n", resource.GoName())
	}
	fmt.Fprintf(w, "	}\n")
	fmt.Fprintf(w, "	return nil\n")
	fmt.Fprintf(w, "}\n")
}

type Resource struct {
	Name       string
	Href       string
	Properties []Property
}

func (r *Resource) IsTopLevelResource() bool {
	return strings.HasPrefix(r.Name, "AWS::")
}

func (r *Resource) Load() error {
	docReader, err := getDoc(r.Href)
	if err != nil {
		return err
	}

	doc, err := goquery.NewDocumentFromReader(docReader)
	if err != nil {
		return err
	}

	// An element with the class 'variablelist' immediately preceeded by an
	// element with the text "Properties" is what we're looking for.
	doc.Find(".variablelist").Each(func(i int, varList *goquery.Selection) {
		if varList.Parent().Find(".titlepage").First().Text() != "Properties" {
			return
		}

		// The variableList contains a definition list. The DT element is the
		// name of the property, the following DD element contains information
		// about it, including the type.
		varList.Find("dl dt").Each(func(i int, dt *goquery.Selection) {
			property := Property{Name: dt.Text()}

			dd := dt.Next()

			property.DocString = dd.Find("p").First().Text()
			property.DocString = regexp.MustCompile("\\s+").ReplaceAllString(property.DocString, " ")
			property.DocString = wordwrap.WrapString(property.DocString, 70)
			/*
				if err != nil {
					fmt.Fprintf(os.Stderr, "%s\n", err)
					return
				}
				property.DocString, err = html2text.FromString(docString)
				if err != nil {
					fmt.Fprintf(os.Stderr, "%s\n", err)
					return
				}
			*/

			// Somewhere inside the
 element there is a span that starts
			// with `Type: ` which is our type. Grab it along with the href
			// from an anchor (for complex types, this tells us which type it
			// refers to)
			dd.Find("span").Each(func(j int, span *goquery.Selection) {
				if span.Text() == "Type" {
					property.Type = span.Parent().Text()
					property.Type = strings.TrimPrefix(property.Type, "Type: ")
					property.Type = regexp.MustCompile("\\s+").ReplaceAllString(property.Type, " ")
					property.Type = strings.TrimSuffix(property.Type, ".")

					span.Parent().Find("a").Each(func(j int, a *goquery.Selection) {
						property.TypeHref, _ = a.Attr("href")
					})
				}
			})
			r.Properties = append(r.Properties, property)
		})
	})
	return nil
}

func (r *Resource) GoName() string {
	rv := r.Name
	rv = strings.TrimPrefix(rv, "AWS CloudFormation ")
	rv = strings.TrimPrefix(rv, "AWS ")
	rv = strings.TrimPrefix(rv, "Amazon ")
	rv = strings.TrimPrefix(rv, "AWS::")
	rv = strings.Replace(rv, "::", "", -1)
	rv = regexp.MustCompile("\\W").ReplaceAllString(rv, "")
	rv = strings.TrimSuffix(rv, "PropertyType")
	rv = strings.TrimSuffix(rv, "Type")

	if rv == "ResourceTags" {
		rv = "ResourceTag"
	}

	// There is an object names AWS::EC2::NetworkInterfaceAttachment and an
	// object named EC2NetworkInterfaceAttachmentType. To avoid a duplicate
	// definition, we have to deconflict them here.
	if r.Name == "EC2 Network Interface Attachment" {
		return "EC2NetworkInterfaceAttachmentType"
	}

	return rv
}

const rootURL = "http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/"

// http://godoc.org/golang.org/x/net/html
// http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-template-resource-type-ref.html
// Topics ->
AWS::AutoScaling::AutoScalingGroup

func getDoc(url string) (io.ReadCloser, error) {
	_, cachePath := path.Split(url)
	cachePath = path.Join("./.scraper-cache", cachePath)
	d, err := os.Open(cachePath)
	if err == nil {
		return d, nil
	}
	if !os.IsNotExist(err) {
		return nil, err
	}
	res, err := http.Get(rootURL + url)
	if err != nil {
		return nil, err
	}

	d, err = os.Create(cachePath)
	if err != nil {
		return nil, err
	}
	io.Copy(d, res.Body)
	d.Close()

	return os.Open(cachePath)
}

type Property struct {
	Name      string
	Type      string
	TypeHref  string
	TypeName  string
	DocString string
}

func (p *Property) GoName() string {
	rv := strings.Title(p.Name)
	rv = strings.Replace(rv, ".", "", -1)
	rv = regexp.MustCompile("[^A-Za-z0-9]").ReplaceAllString(rv, "X")
	return rv
}

func (p *Property) Comment(prefix string) string {
	c := p.DocString
	c = strings.Replace(c, "\n\n: ", ": ", -1)
	c = strings.Replace(c, "\n\n:\n", ": ", -1)
	c = strings.Replace(c, "\n\n", "\n", -1)
	c = strings.Replace(c, "\n", "\n"+prefix, -1)
	return prefix + c + "\n"
}

func (p *Property) GoType(tr *TemplateReference) string {
	if p.TypeHref != "" {
		for _, res := range tr.Resources {
			if res.Href == p.TypeHref {
				p.TypeName = "*" + res.GoName()
			}
		}
		if p.TypeName == "" {
			p.TypeName = "[UNKNOWN " + p.TypeHref + "]"
		}
	}
	if p.TypeName != "" {
		if strings.HasPrefix(p.Type, "A list of") ||
			strings.HasPrefix(p.Type, "List of") ||
			strings.HasPrefix(p.Type, "list of") {
			return p.TypeName + "List"
		}
		// In various places the documentation omit the "list of"
		// when describing types, but include the phrase "list" in the
		// docstring. For example SecurityGroupEgress and SecurityGroupIngress in
		// AWS::EC2::SecurityGroup as "EC2 Security Group Rule" when
		// it should be "list of EC2 Security Group Rule"
		// c.f. http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-security-group.html#cfn-ec2-securitygroup-securitygroupegress
		if strings.HasPrefix(p.DocString, "A list") {
			return p.TypeName + "List"
		}

		if p.Type == "AWS CloudFormation Resource Tags" {
			return "[]ResourceTag"
		}

		return p.TypeName
	}

	switch p.Type {
	case "String":
		return "*StringExpr"
	case "List of strings":
		return "*StringListExpr"
	case "String list":
		return "*StringListExpr"
	case "Boolean":
		return "*BoolExpr"
	case "Integer":
		return "*IntegerExpr"
	case "Number":
		return "*IntegerExpr"
	case "Time stamp":
		return "time.Time"
	}

	if strings.HasPrefix(p.Type, "Number") {
		return "*IntegerExpr"
	}
	if strings.HasPrefix(p.Type, "String") {
		return "*StringExpr"
	}

	return "interface{}"
}

func main() {
	var format = flag.String("format", "go", "How to write the output, either `json` or `go`.")
	var outPath = flag.String("out", "", "The output path")

	flag.Parse()

	tr := TemplateReference{}
	if err := tr.Load(); err != nil {
		log.Fatal(err)
	}

	var out io.Writer
	if *outPath == "-" {
		out = os.Stdout
	} else {
		var err error
		out, err = os.OpenFile(*outPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatal(err)
		}
	}

	switch *format {
	case "go":
		tr.WriteGo(out)
	case "json":
		json.NewEncoder(out).Encode(tr)
	default:
		fmt.Fprintf(os.Stderr, "unrecognised output format: %q", *format)
	}
}
<commit_msg>scraper: don’t break if the scrape cache directory doesn’t exist<commit_after>// This program scrapes the cloudformation documentation to determine the schema
// and produces a go program to the file specified by the `-out` flag.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"regexp"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"github.com/mitchellh/go-wordwrap"
)

// TemplateReference describes the CloudFormation template schema
type TemplateReference struct {
	Resources []*Resource
}

func (tr *TemplateReference) Load() error {
	for _, docURI := range []string{"aws-template-resource-type-ref.html", "aws-product-property-reference.html"} {
		docReader, err := getDoc(docURI)
		if err != nil {
			return err
		}
		doc, err := goquery.NewDocumentFromReader(docReader)
		if err != nil {
			return err
		}
		doc.Find(".highlights li a").Each(func(i int, s *goquery.Selection) {
			name := s.Text()
			name = regexp.MustCompile("\\s+").ReplaceAllString(name, " ")
			href, _ := s.Attr("href")
			tr.Resources = append(tr.Resources, &Resource{Name: name, Href: href})

		})
	}

	for _, resource := range tr.Resources {
		if err := resource.Load(); err != nil {
			return fmt.Errorf("%s: %s", resource.Name, err)
		}
	}
	return nil
}

func (tr *TemplateReference) WriteGo(w io.Writer) {
	fmt.Fprintf(w, "package cloudformation\n")
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "import \"time\"\n")
	fmt.Fprintf(w, "import \"encoding/json\"\n")
	for _, resource := range tr.Resources {
		fmt.Fprintf(w, "\n")
		fmt.Fprintf(w, "// %s represents %s\n", resource.GoName(), resource.Name)
		fmt.Fprintf(w, "//\n")
		fmt.Fprintf(w, "// see %s%s\n", rootURL, resource.Href)
		fmt.Fprintf(w, "type %s struct {\n", resource.GoName())
		for i, p := range resource.Properties {
			if i > 0 {
				fmt.Fprintf(w, "\n")
			}
			fmt.Fprintf(w, "%s", p.Comment("  // "))
			fmt.Fprintf(w, "  %s %s `json:\"%s,omitempty\"`\n", p.GoName(), p.GoType(tr), p.Name)
		}
		fmt.Fprintf(w, "}\n")
		fmt.Fprintf(w, "\n")
		if resource.IsTopLevelResource() {

			fmt.Fprintf(w, "// ResourceType returns %s to implement the ResourceProperties interface\n", resource.Name)
			fmt.Fprintf(w, "func (s %s) ResourceType() string {\n", resource.GoName())
			fmt.Fprintf(w, "	return %q\n", resource.Name)
			fmt.Fprintf(w, "}\n")
			fmt.Fprintf(w, "\n")
		}

		// Cloudformation allows a single object when a list of objects is expected. To
		// handle this we need to generate a *List type. This applies mostly only to
		// non-top-level objects. (The only exception is AWS::Route53::RecordSet which is
		// both a top level resource and a child element of AWS::Route53::RecordSetGroup
		if !resource.IsTopLevelResource() || resource.GoName() == "Route53RecordSet" {
			fmt.Fprintf(w, "// %sList represents a list of %s\n", resource.GoName(), resource.GoName())
			fmt.Fprintf(w, "type %sList []%s\n", resource.GoName(), resource.GoName())
			fmt.Fprintf(w, "\n")
			fmt.Fprintf(w, "// UnmarshalJSON sets the object from the provided JSON representation\n")
			fmt.Fprintf(w, "func (l *%sList) UnmarshalJSON(buf []byte) error {\n", resource.GoName())
			fmt.Fprintf(w, "	// Cloudformation allows a single object when a list of objects is expected\n")
			fmt.Fprintf(w, "	item := %s{}\n", resource.GoName())
			fmt.Fprintf(w, "	if err := json.Unmarshal(buf, &item); err == nil {\n")
			fmt.Fprintf(w, "		*l = %sList{item}\n", resource.GoName())
			fmt.Fprintf(w, "		return nil\n")
			fmt.Fprintf(w, "	}\n")
			fmt.Fprintf(w, "	list := []%s{}\n", resource.GoName())
			fmt.Fprintf(w, "	err := json.Unmarshal(buf, &list)\n")
			fmt.Fprintf(w, "	if err == nil {\n")
			fmt.Fprintf(w, "		*l = %sList(list)\n", resource.GoName())
			fmt.Fprintf(w, "		return nil\n")
			fmt.Fprintf(w, "	}\n")
			fmt.Fprintf(w, "	return err\n")
			fmt.Fprintf(w, "}\n")
			fmt.Fprintf(w, "\n")
		}
	}

	fmt.Fprintf(w, "// NewResourceByType returns a new resource object correspoding with the provided type\n")
	fmt.Fprintf(w, "func NewResourceByType(typeName string) ResourceProperties {\n")
	fmt.Fprintf(w, "	switch typeName {\n")

	for _, resource := range tr.Resources {
		if !resource.IsTopLevelResource() {
			continue
		}
		fmt.Fprintf(w, "		case %q:\n", resource.Name)
		fmt.Fprintf(w, "			return &%s{}\n", resource.GoName())
	}
	fmt.Fprintf(w, "	}\n")
	fmt.Fprintf(w, "	return nil\n")
	fmt.Fprintf(w, "}\n")
}

type Resource struct {
	Name       string
	Href       string
	Properties []Property
}

func (r *Resource) IsTopLevelResource() bool {
	return strings.HasPrefix(r.Name, "AWS::")
}

func (r *Resource) Load() error {
	docReader, err := getDoc(r.Href)
	if err != nil {
		return err
	}

	doc, err := goquery.NewDocumentFromReader(docReader)
	if err != nil {
		return err
	}

	// An element with the class 'variablelist' immediately preceeded by an
	// element with the text "Properties" is what we're looking for.
	doc.Find(".variablelist").Each(func(i int, varList *goquery.Selection) {
		if varList.Parent().Find(".titlepage").First().Text() != "Properties" {
			return
		}

		// The variableList contains a definition list. The DT element is the
		// name of the property, the following DD element contains information
		// about it, including the type.
		varList.Find("dl dt").Each(func(i int, dt *goquery.Selection) {
			property := Property{Name: dt.Text()}

			dd := dt.Next()

			property.DocString = dd.Find("p").First().Text()
			property.DocString = regexp.MustCompile("\\s+").ReplaceAllString(property.DocString, " ")
			property.DocString = wordwrap.WrapString(property.DocString, 70)
			/*
				if err != nil {
					fmt.Fprintf(os.Stderr, "%s\n", err)
					return
				}
				property.DocString, err = html2text.FromString(docString)
				if err != nil {
					fmt.Fprintf(os.Stderr, "%s\n", err)
					return
				}
			*/

			// Somewhere inside the
 element there is a span that starts
			// with `Type: ` which is our type. Grab it along with the href
			// from an anchor (for complex types, this tells us which type it
			// refers to)
			dd.Find("span").Each(func(j int, span *goquery.Selection) {
				if span.Text() == "Type" {
					property.Type = span.Parent().Text()
					property.Type = strings.TrimPrefix(property.Type, "Type: ")
					property.Type = regexp.MustCompile("\\s+").ReplaceAllString(property.Type, " ")
					property.Type = strings.TrimSuffix(property.Type, ".")

					span.Parent().Find("a").Each(func(j int, a *goquery.Selection) {
						property.TypeHref, _ = a.Attr("href")
					})
				}
			})
			r.Properties = append(r.Properties, property)
		})
	})
	return nil
}

func (r *Resource) GoName() string {
	rv := r.Name
	rv = strings.TrimPrefix(rv, "AWS CloudFormation ")
	rv = strings.TrimPrefix(rv, "AWS ")
	rv = strings.TrimPrefix(rv, "Amazon ")
	rv = strings.TrimPrefix(rv, "AWS::")
	rv = strings.Replace(rv, "::", "", -1)
	rv = regexp.MustCompile("\\W").ReplaceAllString(rv, "")
	rv = strings.TrimSuffix(rv, "PropertyType")
	rv = strings.TrimSuffix(rv, "Type")

	if rv == "ResourceTags" {
		rv = "ResourceTag"
	}

	// There is an object names AWS::EC2::NetworkInterfaceAttachment and an
	// object named EC2NetworkInterfaceAttachmentType. To avoid a duplicate
	// definition, we have to deconflict them here.
	if r.Name == "EC2 Network Interface Attachment" {
		return "EC2NetworkInterfaceAttachmentType"
	}

	return rv
}

const rootURL = "http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/"

// http://godoc.org/golang.org/x/net/html
// http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-template-resource-type-ref.html
// Topics ->
AWS::AutoScaling::AutoScalingGroup

func getDoc(url string) (io.ReadCloser, error) {
	_, cachePath := path.Split(url)
	cachePath = path.Join("./.scraper-cache", cachePath)
	os.MkdirAll(path.Dir(cachePath), 0755)

	d, err := os.Open(cachePath)
	if err == nil {
		return d, nil
	}
	if !os.IsNotExist(err) {
		return nil, err
	}
	res, err := http.Get(rootURL + url)
	if err != nil {
		return nil, err
	}

	d, err = os.Create(cachePath)
	if err != nil {
		return nil, err
	}
	io.Copy(d, res.Body)
	d.Close()

	return os.Open(cachePath)
}

type Property struct {
	Name      string
	Type      string
	TypeHref  string
	TypeName  string
	DocString string
}

func (p *Property) GoName() string {
	rv := strings.Title(p.Name)
	rv = strings.Replace(rv, ".", "", -1)
	rv = regexp.MustCompile("[^A-Za-z0-9]").ReplaceAllString(rv, "X")
	return rv
}

func (p *Property) Comment(prefix string) string {
	c := p.DocString
	c = strings.Replace(c, "\n\n: ", ": ", -1)
	c = strings.Replace(c, "\n\n:\n", ": ", -1)
	c = strings.Replace(c, "\n\n", "\n", -1)
	c = strings.Replace(c, "\n", "\n"+prefix, -1)
	return prefix + c + "\n"
}

func (p *Property) GoType(tr *TemplateReference) string {
	if p.TypeHref != "" {
		for _, res := range tr.Resources {
			if res.Href == p.TypeHref {
				p.TypeName = "*" + res.GoName()
			}
		}
		if p.TypeName == "" {
			p.TypeName = "[UNKNOWN " + p.TypeHref + "]"
		}
	}
	if p.TypeName != "" {
		if strings.HasPrefix(p.Type, "A list of") ||
			strings.HasPrefix(p.Type, "List of") ||
			strings.HasPrefix(p.Type, "list of") {
			return p.TypeName + "List"
		}
		// In various places the documentation omit the "list of"
		// when describing types, but include the phrase "list" in the
		// docstring. For example SecurityGroupEgress and SecurityGroupIngress in
		// AWS::EC2::SecurityGroup as "EC2 Security Group Rule" when
		// it should be "list of EC2 Security Group Rule"
		// c.f. http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-security-group.html#cfn-ec2-securitygroup-securitygroupegress
		if strings.HasPrefix(p.DocString, "A list") {
			return p.TypeName + "List"
		}

		if p.Type == "AWS CloudFormation Resource Tags" {
			return "[]ResourceTag"
		}

		return p.TypeName
	}

	switch p.Type {
	case "String":
		return "*StringExpr"
	case "List of strings":
		return "*StringListExpr"
	case "String list":
		return "*StringListExpr"
	case "Boolean":
		return "*BoolExpr"
	case "Integer":
		return "*IntegerExpr"
	case "Number":
		return "*IntegerExpr"
	case "Time stamp":
		return "time.Time"
	}

	if strings.HasPrefix(p.Type, "Number") {
		return "*IntegerExpr"
	}
	if strings.HasPrefix(p.Type, "String") {
		return "*StringExpr"
	}

	return "interface{}"
}

func main() {
	var format = flag.String("format", "go", "How to write the output, either `json` or `go`.")
	var outPath = flag.String("out", "", "The output path")

	flag.Parse()

	tr := TemplateReference{}
	if err := tr.Load(); err != nil {
		log.Fatal(err)
	}

	var out io.Writer
	if *outPath == "-" {
		out = os.Stdout
	} else {
		var err error
		out, err = os.OpenFile(*outPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatal(err)
		}
	}

	switch *format {
	case "go":
		tr.WriteGo(out)
	case "json":
		json.NewEncoder(out).Encode(tr)
	default:
		fmt.Fprintf(os.Stderr, "unrecognised output format: %q", *format)
	}
}
<|endoftext|>


<commit_before>package mesos_connector

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/Dataman-Cloud/swan/src/manager/framework/event"
	"github.com/Dataman-Cloud/swan/src/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/src/mesosproto/sched"
	"github.com/Dataman-Cloud/swan/src/swancontext"

	"github.com/Sirupsen/logrus"
	"github.com/andygrunwald/megos"
	"github.com/golang/protobuf/proto"
	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"
)

var instance *MesosConnector
var once sync.Once

type MesosConnector struct {
	// mesos framework related
	ClusterID        string
	Master           string
	client           *MesosHttpClient
	lastHearBeatTime time.Time

	MesosCallChan chan *sched.Call

	// TODO make sure this chan doesn't explode
	MesosEventChan chan *event.MesosEvent
	Framework      *mesos.FrameworkInfo
}

func NewMesosConnector() *MesosConnector {
	return Instance() // call initialize method
}

func Instance() *MesosConnector {
	once.Do(
		func() {
			instance = &MesosConnector{
				MesosEventChan: make(chan *event.MesosEvent, 1024), // make this unbound in future
				MesosCallChan:  make(chan *sched.Call, 1024),
			}
		})

	return instance
}

func (s *MesosConnector) subscribe(ctx context.Context, mesosFailureChan chan error) {
	logrus.Infof("Subscribe with mesos master %s", s.Master)
	call := &sched.Call{
		Type: sched.Call_SUBSCRIBE.Enum(),
		Subscribe: &sched.Call_Subscribe{
			FrameworkInfo: s.Framework,
		},
	}

	if s.Framework.Id != nil {
		call.FrameworkId = &mesos.FrameworkID{
			Value: proto.String(s.Framework.Id.GetValue()),
		}
	}

	resp, err := s.Send(call)
	if err != nil {
		mesosFailureChan <- err
	}

	// http might now be the default transport in future release
	if resp.StatusCode != http.StatusOK {
		mesosFailureChan <- fmt.Errorf("Subscribe with unexpected response status: %d", resp.StatusCode)
	}

	logrus.Info(s.client.StreamID)
	go s.handleEvents(ctx, resp, mesosFailureChan)
}

func (s *MesosConnector) handleEvents(ctx context.Context, resp *http.Response, mesosFailureChan chan error) {
	defer func() {
		resp.Body.Close()
	}()

	r := NewReader(resp.Body)
	dec := json.NewDecoder(r)

	for {
		select {
		case <-ctx.Done():
			logrus.Infof("handleEvents cancelled %s", ctx.Err())
			return
		default:
			event := new(sched.Event)
			if err := dec.Decode(event); err != nil {
				logrus.Errorf("Deocde event failed: %s", err)
				mesosFailureChan <- err
			}

			switch event.GetType() {
			case sched.Event_SUBSCRIBED:
				s.addEvent(sched.Event_SUBSCRIBED, event)
			case sched.Event_OFFERS:
				s.addEvent(sched.Event_OFFERS, event)
			case sched.Event_RESCIND:
				s.addEvent(sched.Event_RESCIND, event)
			case sched.Event_UPDATE:
				s.addEvent(sched.Event_UPDATE, event)
			case sched.Event_MESSAGE:
				s.addEvent(sched.Event_MESSAGE, event)
			case sched.Event_FAILURE:
				s.addEvent(sched.Event_FAILURE, event)
			case sched.Event_ERROR:
				s.addEvent(sched.Event_ERROR, event)
			case sched.Event_HEARTBEAT:
				s.addEvent(sched.Event_HEARTBEAT, event)
			}
		}
	}
}

func CreateFrameworkInfo() *mesos.FrameworkInfo {
	fw := &mesos.FrameworkInfo{
		User:            proto.String(swancontext.Instance().Config.Scheduler.MesosFrameworkUser),
		Name:            proto.String("swan"),
		FailoverTimeout: proto.Float64(60 * 60 * 24 * 7),
	}

	return fw
}

func getMastersFromZK(zkPath string) ([]string, error) {
	masterInfo := new(mesos.MasterInfo)

	connUrl := zkPath
	if !strings.HasPrefix(connUrl, "zk://") {
		connUrl = fmt.Sprintf("zk://%s", zkPath)
	}
	url, err := url.Parse(connUrl)
	if err != nil {
		return nil, err
	}

	conn, _, err := zk.Connect(strings.Split(url.Host, ","), time.Second)
	defer conn.Close()
	if err != nil {
		return nil, fmt.Errorf("Couldn't connect to zookeeper:%s", err.Error())
	}

	// find mesos master
	children, _, err := conn.Children(url.Path)
	if err != nil {
		return nil, fmt.Errorf("Couldn't connect to zookeeper:%s", err.Error())
	}

	masters := make([]string, 0)
	for _, node := range children {
		if strings.HasPrefix(node, "json.info") {
			data, _, _ := conn.Get(url.Path + "/" + node)
			err := json.Unmarshal(data, masterInfo)
			if err != nil {
				return nil, fmt.Errorf("Unmarshal error: %s", err.Error())
			}
			masters = append(masters, fmt.Sprintf("%s:%d", *masterInfo.GetAddress().Ip, *masterInfo.GetAddress().Port))
		}
	}

	logrus.Info("Find mesos masters: ", masters)
	return masters, nil
}

func stateFromMasters(masters []string) (*megos.State, error) {
	masterUrls := make([]*url.URL, 0)
	for _, master := range masters {
		masterUrl, _ := url.Parse(fmt.Sprintf("http://%s", master))
		masterUrls = append(masterUrls, masterUrl)
	}

	mesos := megos.NewClient(masterUrls, nil)
	return mesos.GetStateFromCluster()
}

func (s *MesosConnector) Send(call *sched.Call) (*http.Response, error) {
	payload, err := proto.Marshal(call)
	if err != nil {
		return nil, err
	}
	return s.client.Send(payload)
}

func (s *MesosConnector) addEvent(eventType sched.Event_Type, e *sched.Event) {
	s.MesosEventChan <- &event.MesosEvent{EventType: eventType, Event: e}
}

func (s *MesosConnector) Start(ctx context.Context, mesosFailureChan chan error) {
	var err error
	masters, err := getMastersFromZK(swancontext.Instance().Config.Scheduler.ZkPath)
	if err != nil {
		logrus.Error(err)
		return
	}
	state, err := stateFromMasters(masters)
	if err != nil {
		logrus.Errorf("%s Check your mesos master configuration", err)
		mesosFailureChan <- err
		return
	}

	s.Master = state.Leader
	s.client = NewHTTPClient(state.Leader, "/api/v1/scheduler")

	s.ClusterID = state.Cluster
	if s.ClusterID == "" {
		s.ClusterID = "cluster"
	}

	r, _ := regexp.Compile("([\\-\\.\\$\\*\\+\\?\\{\\}\\(\\)\\[\\]\\|]+)")
	match := r.MatchString(s.ClusterID)
	if match {
		logrus.Warnf(`Swan do not work with mesos cluster name(%s) with special characters "-.$*+?{}()[]|".`, s.ClusterID)
		s.ClusterID = r.ReplaceAllString(s.ClusterID, "")
		logrus.Infof("Swan acceptable cluster name: %s", s.ClusterID)
	}

	s.subscribe(ctx, mesosFailureChan)

	for {
		select {
		case <-ctx.Done():
			logrus.Errorf("mesosConnector got signal %s", ctx.Err())
			return
		case call := <-s.MesosCallChan:
			logrus.WithFields(logrus.Fields{"sending-call": sched.Call_Type_name[int32(*call.Type)]}).Debugf("%+v", call)
			resp, err := s.Send(call)
			if err != nil {
				logrus.Errorf("%s", err)
				mesosFailureChan <- err
			}
			if resp.StatusCode != 202 {
				logrus.Infof("send response not 202 but %d", resp.StatusCode)
				mesosFailureChan <- errors.New("http got respose not 202")
			}
		}
	}
}
<commit_msg>add log<commit_after>package mesos_connector

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/Dataman-Cloud/swan/src/manager/framework/event"
	"github.com/Dataman-Cloud/swan/src/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/src/mesosproto/sched"
	"github.com/Dataman-Cloud/swan/src/swancontext"

	"github.com/Sirupsen/logrus"
	"github.com/andygrunwald/megos"
	"github.com/golang/protobuf/proto"
	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"
)

var instance *MesosConnector
var once sync.Once

type MesosConnector struct {
	// mesos framework related
	ClusterID        string
	Master           string
	client           *MesosHttpClient
	lastHearBeatTime time.Time

	MesosCallChan chan *sched.Call

	// TODO make sure this chan doesn't explode
	MesosEventChan chan *event.MesosEvent
	Framework      *mesos.FrameworkInfo
}

func NewMesosConnector() *MesosConnector {
	return Instance() // call initialize method
}

func Instance() *MesosConnector {
	once.Do(
		func() {
			instance = &MesosConnector{
				MesosEventChan: make(chan *event.MesosEvent, 1024), // make this unbound in future
				MesosCallChan:  make(chan *sched.Call, 1024),
			}
		})

	return instance
}

func (s *MesosConnector) subscribe(ctx context.Context, mesosFailureChan chan error) {
	logrus.Infof("Subscribe with mesos master %s", s.Master)
	call := &sched.Call{
		Type: sched.Call_SUBSCRIBE.Enum(),
		Subscribe: &sched.Call_Subscribe{
			FrameworkInfo: s.Framework,
		},
	}

	if s.Framework.Id != nil {
		call.FrameworkId = &mesos.FrameworkID{
			Value: proto.String(s.Framework.Id.GetValue()),
		}
	}

	resp, err := s.Send(call)
	if err != nil {
		mesosFailureChan <- err
	}

	// http might now be the default transport in future release
	if resp.StatusCode != http.StatusOK {
		mesosFailureChan <- fmt.Errorf("Subscribe with unexpected response status: %d", resp.StatusCode)
	}

	go s.handleEvents(ctx, resp, mesosFailureChan)
}

func (s *MesosConnector) handleEvents(ctx context.Context, resp *http.Response, mesosFailureChan chan error) {
	defer func() {
		resp.Body.Close()
	}()

	r := NewReader(resp.Body)
	dec := json.NewDecoder(r)

	for {
		select {
		case <-ctx.Done():
			logrus.Infof("handleEvents cancelled %s", ctx.Err())
			return
		default:
			event := new(sched.Event)
			if err := dec.Decode(event); err != nil {
				logrus.Errorf("Deocde event failed: %s", err)
				mesosFailureChan <- err
			}

			switch event.GetType() {
			case sched.Event_SUBSCRIBED:
				logrus.Infof("Subscribed successful with ID %s", event.GetSubscribed().FrameworkId.GetValue())
				s.addEvent(sched.Event_SUBSCRIBED, event)
			case sched.Event_OFFERS:
				s.addEvent(sched.Event_OFFERS, event)
			case sched.Event_RESCIND:
				s.addEvent(sched.Event_RESCIND, event)
			case sched.Event_UPDATE:
				s.addEvent(sched.Event_UPDATE, event)
			case sched.Event_MESSAGE:
				s.addEvent(sched.Event_MESSAGE, event)
			case sched.Event_FAILURE:
				s.addEvent(sched.Event_FAILURE, event)
			case sched.Event_ERROR:
				s.addEvent(sched.Event_ERROR, event)
			case sched.Event_HEARTBEAT:
				s.addEvent(sched.Event_HEARTBEAT, event)
			}
		}
	}
}

func CreateFrameworkInfo() *mesos.FrameworkInfo {
	fw := &mesos.FrameworkInfo{
		User:            proto.String(swancontext.Instance().Config.Scheduler.MesosFrameworkUser),
		Name:            proto.String("swan"),
		FailoverTimeout: proto.Float64(60 * 60 * 24 * 7),
	}

	return fw
}

func getMastersFromZK(zkPath string) ([]string, error) {
	masterInfo := new(mesos.MasterInfo)

	connUrl := zkPath
	if !strings.HasPrefix(connUrl, "zk://") {
		connUrl = fmt.Sprintf("zk://%s", zkPath)
	}
	url, err := url.Parse(connUrl)
	if err != nil {
		return nil, err
	}

	conn, _, err := zk.Connect(strings.Split(url.Host, ","), time.Second)
	defer conn.Close()
	if err != nil {
		return nil, fmt.Errorf("Couldn't connect to zookeeper:%s", err.Error())
	}

	// find mesos master
	children, _, err := conn.Children(url.Path)
	if err != nil {
		return nil, fmt.Errorf("Couldn't connect to zookeeper:%s", err.Error())
	}

	masters := make([]string, 0)
	for _, node := range children {
		if strings.HasPrefix(node, "json.info") {
			data, _, _ := conn.Get(url.Path + "/" + node)
			err := json.Unmarshal(data, masterInfo)
			if err != nil {
				return nil, fmt.Errorf("Unmarshal error: %s", err.Error())
			}
			masters = append(masters, fmt.Sprintf("%s:%d", *masterInfo.GetAddress().Ip, *masterInfo.GetAddress().Port))
		}
	}

	logrus.Info("Find mesos masters: ", masters)
	return masters, nil
}

func stateFromMasters(masters []string) (*megos.State, error) {
	masterUrls := make([]*url.URL, 0)
	for _, master := range masters {
		masterUrl, _ := url.Parse(fmt.Sprintf("http://%s", master))
		masterUrls = append(masterUrls, masterUrl)
	}

	mesos := megos.NewClient(masterUrls, nil)
	return mesos.GetStateFromCluster()
}

func (s *MesosConnector) Send(call *sched.Call) (*http.Response, error) {
	payload, err := proto.Marshal(call)
	if err != nil {
		return nil, err
	}
	return s.client.Send(payload)
}

func (s *MesosConnector) addEvent(eventType sched.Event_Type, e *sched.Event) {
	s.MesosEventChan <- &event.MesosEvent{EventType: eventType, Event: e}
}

func (s *MesosConnector) Start(ctx context.Context, mesosFailureChan chan error) {
	var err error
	masters, err := getMastersFromZK(swancontext.Instance().Config.Scheduler.ZkPath)
	if err != nil {
		logrus.Error(err)
		return
	}
	state, err := stateFromMasters(masters)
	if err != nil {
		logrus.Errorf("%s Check your mesos master configuration", err)
		mesosFailureChan <- err
		return
	}

	s.Master = state.Leader
	s.client = NewHTTPClient(state.Leader, "/api/v1/scheduler")

	s.ClusterID = state.Cluster
	if s.ClusterID == "" {
		s.ClusterID = "cluster"
	}

	r, _ := regexp.Compile("([\\-\\.\\$\\*\\+\\?\\{\\}\\(\\)\\[\\]\\|]+)")
	match := r.MatchString(s.ClusterID)
	if match {
		logrus.Warnf(`Swan do not work with mesos cluster name(%s) with special characters "-.$*+?{}()[]|".`, s.ClusterID)
		s.ClusterID = r.ReplaceAllString(s.ClusterID, "")
		logrus.Infof("Swan acceptable cluster name: %s", s.ClusterID)
	}

	s.subscribe(ctx, mesosFailureChan)

	for {
		select {
		case <-ctx.Done():
			logrus.Errorf("mesosConnector got signal %s", ctx.Err())
			return
		case call := <-s.MesosCallChan:
			logrus.WithFields(logrus.Fields{"sending-call": sched.Call_Type_name[int32(*call.Type)]}).Debugf("%+v", call)
			resp, err := s.Send(call)
			if err != nil {
				logrus.Errorf("%s", err)
				mesosFailureChan <- err
			}
			if resp.StatusCode != 202 {
				logrus.Infof("send response not 202 but %d", resp.StatusCode)
				mesosFailureChan <- errors.New("http got respose not 202")
			}
		}
	}
}
<|endoftext|>


<commit_before>package git_pipeline_test

import (
	"os"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/cloudfoundry-incubator/garden/client"
	"github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/testflight/bosh"
	"github.com/concourse/testflight/gitserver"
	"github.com/concourse/testflight/guidserver"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
	"time"
)

const helperRootfs = "docker:///concourse/testflight-helper"

var (
	gardenClient garden.Client

	gitServer *gitserver.Server

	successGitServer  *gitserver.Server
	failureGitServer  *gitserver.Server
	noUpdateGitServer *gitserver.Server
)

type GardenLinuxDeploymentData struct {
	GardenLinuxVersion string
}

type GitPipelineTemplate struct {
	GitServers struct {
		Origin   string
		Success  string
		Failure  string
		NoUpdate string
	}

	GuidServerCurlCommand string

	TestflightHelperImage string

	GardenLinuxDeploymentData
}

var _ = BeforeSuite(func() {
	gardenLinuxVersion := os.Getenv("GARDEN_LINUX_VERSION")
	Ω(gardenLinuxVersion).ShouldNot(BeEmpty(), "must set $GARDEN_LINUX_VERSION")

	bosh.DeleteDeployment("garden-testflight")
	bosh.DeleteDeployment("concourse-testflight")

	gardenLinuxDeploymentData := GardenLinuxDeploymentData{
		GardenLinuxVersion: gardenLinuxVersion,
	}

	bosh.Deploy("garden.yml", gardenLinuxDeploymentData)

	gardenClient = client.New(connection.New("tcp", "10.244.16.2:7777"))
	Eventually(gardenClient.Ping, 10*time.Second).ShouldNot(HaveOccurred())

	guidserver.Start(helperRootfs, gardenClient)

	gitServer = gitserver.Start(helperRootfs, gardenClient)
	successGitServer = gitserver.Start(helperRootfs, gardenClient)
	failureGitServer = gitserver.Start(helperRootfs, gardenClient)
	noUpdateGitServer = gitserver.Start(helperRootfs, gardenClient)

	templateData := GitPipelineTemplate{
		GardenLinuxDeploymentData: gardenLinuxDeploymentData,
	}

	templateData.GitServers.Origin = gitServer.URI()
	templateData.GitServers.Success = successGitServer.URI()
	templateData.GitServers.Failure = failureGitServer.URI()
	templateData.GitServers.NoUpdate = noUpdateGitServer.URI()

	templateData.TestflightHelperImage = helperRootfs
	templateData.GuidServerCurlCommand = guidserver.CurlCommand()

	bosh.Deploy("deployment.yml.tmpl", templateData)
})

var _ = AfterSuite(func() {
	gitServer.Stop()
	successGitServer.Stop()
	failureGitServer.Stop()
	noUpdateGitServer.Stop()

	guidserver.Stop(gardenClient)
})

func TestGitPipeline(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Git Pipeline Suite")
}
<commit_msg>.tmpl<commit_after>package git_pipeline_test

import (
	"os"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/cloudfoundry-incubator/garden/client"
	"github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/testflight/bosh"
	"github.com/concourse/testflight/gitserver"
	"github.com/concourse/testflight/guidserver"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
	"time"
)

const helperRootfs = "docker:///concourse/testflight-helper"

var (
	gardenClient garden.Client

	gitServer *gitserver.Server

	successGitServer  *gitserver.Server
	failureGitServer  *gitserver.Server
	noUpdateGitServer *gitserver.Server
)

type GardenLinuxDeploymentData struct {
	GardenLinuxVersion string
}

type GitPipelineTemplate struct {
	GitServers struct {
		Origin   string
		Success  string
		Failure  string
		NoUpdate string
	}

	GuidServerCurlCommand string

	TestflightHelperImage string

	GardenLinuxDeploymentData
}

var _ = BeforeSuite(func() {
	gardenLinuxVersion := os.Getenv("GARDEN_LINUX_VERSION")
	Ω(gardenLinuxVersion).ShouldNot(BeEmpty(), "must set $GARDEN_LINUX_VERSION")

	bosh.DeleteDeployment("garden-testflight")
	bosh.DeleteDeployment("concourse-testflight")

	gardenLinuxDeploymentData := GardenLinuxDeploymentData{
		GardenLinuxVersion: gardenLinuxVersion,
	}

	bosh.Deploy("garden.yml.tmpl", gardenLinuxDeploymentData)

	gardenClient = client.New(connection.New("tcp", "10.244.16.2:7777"))
	Eventually(gardenClient.Ping, 10*time.Second).ShouldNot(HaveOccurred())

	guidserver.Start(helperRootfs, gardenClient)

	gitServer = gitserver.Start(helperRootfs, gardenClient)
	successGitServer = gitserver.Start(helperRootfs, gardenClient)
	failureGitServer = gitserver.Start(helperRootfs, gardenClient)
	noUpdateGitServer = gitserver.Start(helperRootfs, gardenClient)

	templateData := GitPipelineTemplate{
		GardenLinuxDeploymentData: gardenLinuxDeploymentData,
	}

	templateData.GitServers.Origin = gitServer.URI()
	templateData.GitServers.Success = successGitServer.URI()
	templateData.GitServers.Failure = failureGitServer.URI()
	templateData.GitServers.NoUpdate = noUpdateGitServer.URI()

	templateData.TestflightHelperImage = helperRootfs
	templateData.GuidServerCurlCommand = guidserver.CurlCommand()

	bosh.Deploy("deployment.yml.tmpl", templateData)
})

var _ = AfterSuite(func() {
	gitServer.Stop()
	successGitServer.Stop()
	failureGitServer.Stop()
	noUpdateGitServer.Stop()

	guidserver.Stop(gardenClient)
})

func TestGitPipeline(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Git Pipeline Suite")
}
<|endoftext|>


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

// Package configure provides a very simple gnu configure/make style configure
// script generating a simple Makefile and go file containing all the configured
// variables.
package configure

import (
	"bytes"
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"os"
	"path"
	"regexp"
	"sort"
	"strings"
	"runtime"
)

// Options contains all the standard configure options to specify various
// directories. Use NewOptions to create an instance of this type with the
// common default values for each variable.
type Options struct {
	Prefix        string `long:"prefix" description:"install architecture-independent files in PREFIX"`
	ExecPrefix    string `long:"execprefix" description:"install architecture-dependent files in EPREFIX"`
	BinDir        string `long:"bindir" description:"user executables"`
	LibExecDir    string `long:"libexecdir" description:"program executables"`
	SysConfDir    string `long:"sysconfdir" description:"read-only single-machine data"`
	LibDir        string `long:"libdir" description:"program executables"`
	DataRootDir   string `long:"datarootdir" description:"read-only arch.-independent data root"`
	DataDir       string `long:"datadir" description:"read-only arc.-independent data"`
	ManDir        string `long:"mandir" description:"man documentation"`
}

// NewOptions creates a new Options with common default values.
func NewOptions() *Options {
	return &Options{
		Prefix:        "/usr/local",
		ExecPrefix:    "${prefix}",
		BinDir:        "${execprefix}/bin",
		LibExecDir:    "${execprefix}/libexec",
		LibDir:        "${execprefix}/lib",
		SysConfDir:    "${prefix}/etc",
		DataRootDir:   "${prefix}/share",
		DataDir:       "${datarootdir}",
		ManDir:        "${datarootdir}/man",
	}
}

// Package is the package name in which the GoConfig file will be written
var Package = "main"

// Makefile is the filename of the makefile that will be generated
var Makefile = "go.make"

// GoConfig is the filename of the go file that will be generated containing
// all the variable values.
var GoConfig = "appconfig"

// GoConfigVariable is the name of the variable inside the GoConfig file
// containing all the variable values.
var GoConfigVariable = "AppConfig"

// Target is the executable name to build. If left empty, the name is deduced
// from the directory (similar to what go does)
var Target = ""

// Version is the application version
var Version []int = []int{0, 1}

type expandStringPart struct {
	Value      string
	IsVariable bool
}

func (x *expandStringPart) expand(m map[string]*expandString) (string, []string) {
	if x.IsVariable {
		s, ok := m[x.Value]

		if !ok {
			return "", nil
		} else {
			ret := s.expand(m)
			rets := make([]string, len(s.dependencies), len(s.dependencies)+1)

			copy(rets, s.dependencies)

			return ret, append(rets, x.Value)
		}
	}

	return x.Value, nil
}

type expandString struct {
	Name  string
	Parts []expandStringPart

	dependencies []string
	value        string
	hasExpanded  bool
}

func (x *expandString) dependsOn(name string) bool {
	i := sort.SearchStrings(x.dependencies, name)

	return i < len(x.dependencies) && x.dependencies[i] == name
}

func (x *expandString) expand(m map[string]*expandString) string {
	if !x.hasExpanded {
		// Prevent infinite loop by circular dependencies
		x.hasExpanded = true
		buf := bytes.Buffer{}

		for _, v := range x.Parts {
			s, deps := v.expand(m)
			buf.WriteString(s)

			x.dependencies = append(x.dependencies, deps...)
		}

		sort.Strings(x.dependencies)
		x.value = buf.String()
	}

	return x.value
}

// Config represents the current configuration. See Configure for more
// information.
type Config struct {
	*flags.Parser

	values []*flags.Option
	valuesMap   map[string]*flags.Option
	expanded map[string]*expandString
}

func eachGroup(g *flags.Group, f func(g *flags.Group)) {
	f(g)

	for _, gg := range g.Groups() {
		eachGroup(gg, f)
	}
}

func (x *Config) extract() ([]*flags.Option, map[string]*flags.Option) {
	valuesmap := make(map[string]*flags.Option)
	var values []*flags.Option

	eachGroup(x.Parser.Command.Group, func(g *flags.Group) {
		for _, option := range g.Options() {
			if len(option.LongName) > 0 {
				valuesmap[option.LongName] = option
				values = append(values, option)
			}
		}
	})

	return values, valuesmap
}

func (x *Config) expand() map[string]*expandString {
	ret := make(map[string]*expandString)

	r, _ := regexp.Compile(`\$\{[^}]*\}`)

	for name, opt := range x.valuesMap {
		es := expandString{
			Name: name,
		}

		// Find all variable references
		s, ok := opt.Value().(string)

		if !ok {
			continue
		}

		matches := r.FindAllStringIndex(s, -1)

		for i, match := range matches {
			var prefix string

			if i == 0 {
				prefix = s[0:match[0]]
			} else {
				prefix = s[matches[i-1][1]:match[0]]
			}

			if len(prefix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: prefix, IsVariable: false})
			}

			varname := s[match[0]+2 : match[1]-1]
			es.Parts = append(es.Parts, expandStringPart{Value: varname, IsVariable: true})
		}

		if len(matches) == 0 {
			es.Parts = append(es.Parts, expandStringPart{Value: s, IsVariable: false})
		} else {
			last := matches[len(matches)-1]
			suffix := s[last[1]:]

			if len(suffix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: suffix, IsVariable: false})
			}
		}

		ret[name] = &es
	}

	for _, val := range ret {
		val.expand(ret)
	}

	return ret
}

// Configure runs the configure process with options as provided by the given
// data variable. If data is nil, the default options will be used
// (see NewOptions). Note that the data provided is simply passed to go-flags.
// For more information on flags parsing, see the documentation of go-flags.
// If GoConfig is not empty, then the go configuration will be written to the
// GoConfig file. Similarly, if Makefile is not empty, the Makefile will be
// written.
func Configure(data interface{}) (*Config, error) {
	if data == nil {
		data = NewOptions()
	}

	parser := flags.NewParser(data, flags.PrintErrors | flags.IgnoreUnknown)

	if _, err := parser.Parse(); err != nil {
		return nil, err
	}

	ret := &Config{
		Parser: parser,
	}

	ret.values, ret.valuesMap = ret.extract()
	ret.expanded = ret.expand()

	if len(GoConfig) != 0 {
		filename := GoConfig

		if !strings.HasSuffix(filename, ".go") {
			filename += ".go"
		}

		f, err := os.Create(filename)

		if err != nil {
			return nil, err
		}

		ret.WriteGoConfig(f)
		f.Close()
	}

	if len(Makefile) != 0 {
		f, err := os.Create(Makefile)

		if err != nil {
			return nil, err
		}

		ret.WriteMakefile(f)
		f.Close()

		os.Chmod(Makefile, 0755)

		f, err = os.OpenFile(path.Join(path.Dir(Makefile), "Makefile"),
		                     os.O_CREATE | os.O_EXCL | os.O_WRONLY,
		                     0644)

		if err == nil {
			fmt.Fprintf(f, "include %s\n", path.Base(Makefile))
			f.Close()
		}
	}

	return ret, nil
}

// Expand expands the variable value indicated by name
func (x *Config) Expand(name string) string {
	return x.expanded[name].expand(x.expanded)
}

// WriteGoConfig writes the go configuration file containing all the variable
// values to the given writer. Note that it will write a package line if
// the Package variable is not empty. The GoConfigVariable name will
// be used as the variable name for the configuration.
func (x *Config) WriteGoConfig(writer io.Writer) {
	if len(Package) > 0 {
		fmt.Fprintf(writer, "package %v\n\n", Package)
	}

	fmt.Fprintf(writer, "var %s = struct {\n", GoConfigVariable)
	values := make([]string, 0)

	variables := make([]string, len(x.values))

	// Write all options
	for i, opt := range x.values {
		variables[i] = opt.LongName
	}

	sort.Strings(variables)

	for i, name := range variables {
		if i != 0 {
			io.WriteString(writer, "\n")
		}

		option := x.valuesMap[name]
		val := option.Value()

		fmt.Fprintf(writer, "\t// %s\n", option.Description)
		fmt.Fprintf(writer, "\t%v %T\n", strings.Title(name), val)

		var value string

		if _, ok := x.expanded[option.LongName]; ok {
			value = fmt.Sprintf("%#v", x.Expand(option.LongName))
		} else {
			value = fmt.Sprintf("%#v", val)
		}

		values = append(values, value)
	}

	if len(variables) > 0 {
		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "\t// Application version\n")
	io.WriteString(writer, "\tVersion []int\n")
	fmt.Fprintln(writer, "}{")

	for _, v := range values {
		fmt.Fprintf(writer, "\t%v,\n", v)
	}

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ", ")
		} else {
			io.WriteString(writer, "\t[]int{")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	fmt.Fprintln(writer, "},")
	fmt.Fprintln(writer, "}")
}

// WriteMakefile writes a Makefile for the given parser to the given writer.
// The Makefile contains the common build, clean, distclean, install and
// uninstall rules.
func (x *Config) WriteMakefile(writer io.Writer) {
	// Write a very basic makefile
	io.WriteString(writer, "#!/usr/bin/make -f\n\n")

	vars := make([]*expandString, 0, len(x.expanded))

	for name, v := range x.expanded {
		inserted := false

		// Insert into vars based on dependencies
		for i, vv := range vars {
			if vv.dependsOn(name) {
				tail := make([]*expandString, len(vars)-i)
				copy(tail, vars[i:])

				if i == 0 {
					vars = append([]*expandString{v}, vars...)
				} else {
					vars = append(append(vars[0:i], v), tail...)
				}

				inserted = true
				break
			}
		}

		if !inserted {
			vars = append(vars, v)
		}
	}

	io.WriteString(writer, "# Variables\n")

	for _, v := range vars {
		fmt.Fprintf(writer, "%s ?= ", v.Name)

		for _, part := range v.Parts {
			if part.IsVariable {
				fmt.Fprintf(writer, "$(%s)", part.Value)
			} else {
				fmt.Fprintf(writer, "%s", part.Value)
			}
		}

		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "version ?= ")

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ".")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	io.WriteString(writer, "\n")
	fmt.Fprintf(writer, "major_version = %v\n", Version[0])

	if len(Version) > 1 {
		fmt.Fprintf(writer, "minor_version = %v\n", Version[1])
	}

	if len(Version) > 2 {
		fmt.Fprintf(writer, "micro_version = %v\n", Version[2])
	}

	io.WriteString(writer, "\n")

	target := Target

	if len(target) == 0 {
		pc := make([]uintptr, 3)
		n := runtime.Callers(1, pc)

		me, _ := runtime.FuncForPC(pc[0]).FileLine(pc[0])

		for i := 1; i < n; i++ {
			f := runtime.FuncForPC(pc[i])
			fname, _ := f.FileLine(pc[i])

			if fname != me {
				target = path.Base(path.Dir(fname))
				break
			}
		}
	}

	fmt.Fprintf(writer, "TARGET ?= %s\n", target)

	io.WriteString(writer, "\nSOURCES ?=")
	io.WriteString(writer, "\nSOURCES += $(wildcard *.go)")
	io.WriteString(writer, "\nSOURCES_UNIQUE = $(sort $(SOURCES))")

	io.WriteString(writer, "\n\n")

	io.WriteString(writer, "# Rules\n")
	io.WriteString(writer, "$(TARGET): $(SOURCES_UNIQUE)\n")
	io.WriteString(writer, "\tgo build -o $@\n\n")

	io.WriteString(writer, "clean:\n")
	io.WriteString(writer, "\trm -f $(TARGET)\n\n")

	io.WriteString(writer, "distclean: clean\n\n")

	io.WriteString(writer, "$(TARGET)_installdir ?= $(bindir)\n\n")

	io.WriteString(writer, "install: $(TARGET)\n")
	io.WriteString(writer, "\tmkdir -p $(DESTDIR)$($(TARGET)_installdir) && cp $(TARGET) $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, "uninstall:\n")
	io.WriteString(writer, "\trm -f $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, ".PHONY: install uninstall distclean clean")
}
<commit_msg>Run go fmt<commit_after>// Copyright 2012 Jesse van den Kieboom. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package configure provides a very simple gnu configure/make style configure
// script generating a simple Makefile and go file containing all the configured
// variables.
package configure

import (
	"bytes"
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"os"
	"path"
	"regexp"
	"runtime"
	"sort"
	"strings"
)

// Options contains all the standard configure options to specify various
// directories. Use NewOptions to create an instance of this type with the
// common default values for each variable.
type Options struct {
	Prefix      string `long:"prefix" description:"install architecture-independent files in PREFIX"`
	ExecPrefix  string `long:"execprefix" description:"install architecture-dependent files in EPREFIX"`
	BinDir      string `long:"bindir" description:"user executables"`
	LibExecDir  string `long:"libexecdir" description:"program executables"`
	SysConfDir  string `long:"sysconfdir" description:"read-only single-machine data"`
	LibDir      string `long:"libdir" description:"program executables"`
	DataRootDir string `long:"datarootdir" description:"read-only arch.-independent data root"`
	DataDir     string `long:"datadir" description:"read-only arc.-independent data"`
	ManDir      string `long:"mandir" description:"man documentation"`
}

// NewOptions creates a new Options with common default values.
func NewOptions() *Options {
	return &Options{
		Prefix:      "/usr/local",
		ExecPrefix:  "${prefix}",
		BinDir:      "${execprefix}/bin",
		LibExecDir:  "${execprefix}/libexec",
		LibDir:      "${execprefix}/lib",
		SysConfDir:  "${prefix}/etc",
		DataRootDir: "${prefix}/share",
		DataDir:     "${datarootdir}",
		ManDir:      "${datarootdir}/man",
	}
}

// Package is the package name in which the GoConfig file will be written
var Package = "main"

// Makefile is the filename of the makefile that will be generated
var Makefile = "go.make"

// GoConfig is the filename of the go file that will be generated containing
// all the variable values.
var GoConfig = "appconfig"

// GoConfigVariable is the name of the variable inside the GoConfig file
// containing all the variable values.
var GoConfigVariable = "AppConfig"

// Target is the executable name to build. If left empty, the name is deduced
// from the directory (similar to what go does)
var Target = ""

// Version is the application version
var Version []int = []int{0, 1}

type expandStringPart struct {
	Value      string
	IsVariable bool
}

func (x *expandStringPart) expand(m map[string]*expandString) (string, []string) {
	if x.IsVariable {
		s, ok := m[x.Value]

		if !ok {
			return "", nil
		} else {
			ret := s.expand(m)
			rets := make([]string, len(s.dependencies), len(s.dependencies)+1)

			copy(rets, s.dependencies)

			return ret, append(rets, x.Value)
		}
	}

	return x.Value, nil
}

type expandString struct {
	Name  string
	Parts []expandStringPart

	dependencies []string
	value        string
	hasExpanded  bool
}

func (x *expandString) dependsOn(name string) bool {
	i := sort.SearchStrings(x.dependencies, name)

	return i < len(x.dependencies) && x.dependencies[i] == name
}

func (x *expandString) expand(m map[string]*expandString) string {
	if !x.hasExpanded {
		// Prevent infinite loop by circular dependencies
		x.hasExpanded = true
		buf := bytes.Buffer{}

		for _, v := range x.Parts {
			s, deps := v.expand(m)
			buf.WriteString(s)

			x.dependencies = append(x.dependencies, deps...)
		}

		sort.Strings(x.dependencies)
		x.value = buf.String()
	}

	return x.value
}

// Config represents the current configuration. See Configure for more
// information.
type Config struct {
	*flags.Parser

	values    []*flags.Option
	valuesMap map[string]*flags.Option
	expanded  map[string]*expandString
}

func eachGroup(g *flags.Group, f func(g *flags.Group)) {
	f(g)

	for _, gg := range g.Groups() {
		eachGroup(gg, f)
	}
}

func (x *Config) extract() ([]*flags.Option, map[string]*flags.Option) {
	valuesmap := make(map[string]*flags.Option)
	var values []*flags.Option

	eachGroup(x.Parser.Command.Group, func(g *flags.Group) {
		for _, option := range g.Options() {
			if len(option.LongName) > 0 {
				valuesmap[option.LongName] = option
				values = append(values, option)
			}
		}
	})

	return values, valuesmap
}

func (x *Config) expand() map[string]*expandString {
	ret := make(map[string]*expandString)

	r, _ := regexp.Compile(`\$\{[^}]*\}`)

	for name, opt := range x.valuesMap {
		es := expandString{
			Name: name,
		}

		// Find all variable references
		s, ok := opt.Value().(string)

		if !ok {
			continue
		}

		matches := r.FindAllStringIndex(s, -1)

		for i, match := range matches {
			var prefix string

			if i == 0 {
				prefix = s[0:match[0]]
			} else {
				prefix = s[matches[i-1][1]:match[0]]
			}

			if len(prefix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: prefix, IsVariable: false})
			}

			varname := s[match[0]+2 : match[1]-1]
			es.Parts = append(es.Parts, expandStringPart{Value: varname, IsVariable: true})
		}

		if len(matches) == 0 {
			es.Parts = append(es.Parts, expandStringPart{Value: s, IsVariable: false})
		} else {
			last := matches[len(matches)-1]
			suffix := s[last[1]:]

			if len(suffix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: suffix, IsVariable: false})
			}
		}

		ret[name] = &es
	}

	for _, val := range ret {
		val.expand(ret)
	}

	return ret
}

// Configure runs the configure process with options as provided by the given
// data variable. If data is nil, the default options will be used
// (see NewOptions). Note that the data provided is simply passed to go-flags.
// For more information on flags parsing, see the documentation of go-flags.
// If GoConfig is not empty, then the go configuration will be written to the
// GoConfig file. Similarly, if Makefile is not empty, the Makefile will be
// written.
func Configure(data interface{}) (*Config, error) {
	if data == nil {
		data = NewOptions()
	}

	parser := flags.NewParser(data, flags.PrintErrors|flags.IgnoreUnknown)

	if _, err := parser.Parse(); err != nil {
		return nil, err
	}

	ret := &Config{
		Parser: parser,
	}

	ret.values, ret.valuesMap = ret.extract()
	ret.expanded = ret.expand()

	if len(GoConfig) != 0 {
		filename := GoConfig

		if !strings.HasSuffix(filename, ".go") {
			filename += ".go"
		}

		f, err := os.Create(filename)

		if err != nil {
			return nil, err
		}

		ret.WriteGoConfig(f)
		f.Close()
	}

	if len(Makefile) != 0 {
		f, err := os.Create(Makefile)

		if err != nil {
			return nil, err
		}

		ret.WriteMakefile(f)
		f.Close()

		os.Chmod(Makefile, 0755)

		f, err = os.OpenFile(path.Join(path.Dir(Makefile), "Makefile"),
			os.O_CREATE|os.O_EXCL|os.O_WRONLY,
			0644)

		if err == nil {
			fmt.Fprintf(f, "include %s\n", path.Base(Makefile))
			f.Close()
		}
	}

	return ret, nil
}

// Expand expands the variable value indicated by name
func (x *Config) Expand(name string) string {
	return x.expanded[name].expand(x.expanded)
}

// WriteGoConfig writes the go configuration file containing all the variable
// values to the given writer. Note that it will write a package line if
// the Package variable is not empty. The GoConfigVariable name will
// be used as the variable name for the configuration.
func (x *Config) WriteGoConfig(writer io.Writer) {
	if len(Package) > 0 {
		fmt.Fprintf(writer, "package %v\n\n", Package)
	}

	fmt.Fprintf(writer, "var %s = struct {\n", GoConfigVariable)
	values := make([]string, 0)

	variables := make([]string, len(x.values))

	// Write all options
	for i, opt := range x.values {
		variables[i] = opt.LongName
	}

	sort.Strings(variables)

	for i, name := range variables {
		if i != 0 {
			io.WriteString(writer, "\n")
		}

		option := x.valuesMap[name]
		val := option.Value()

		fmt.Fprintf(writer, "\t// %s\n", option.Description)
		fmt.Fprintf(writer, "\t%v %T\n", strings.Title(name), val)

		var value string

		if _, ok := x.expanded[option.LongName]; ok {
			value = fmt.Sprintf("%#v", x.Expand(option.LongName))
		} else {
			value = fmt.Sprintf("%#v", val)
		}

		values = append(values, value)
	}

	if len(variables) > 0 {
		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "\t// Application version\n")
	io.WriteString(writer, "\tVersion []int\n")
	fmt.Fprintln(writer, "}{")

	for _, v := range values {
		fmt.Fprintf(writer, "\t%v,\n", v)
	}

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ", ")
		} else {
			io.WriteString(writer, "\t[]int{")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	fmt.Fprintln(writer, "},")
	fmt.Fprintln(writer, "}")
}

// WriteMakefile writes a Makefile for the given parser to the given writer.
// The Makefile contains the common build, clean, distclean, install and
// uninstall rules.
func (x *Config) WriteMakefile(writer io.Writer) {
	// Write a very basic makefile
	io.WriteString(writer, "#!/usr/bin/make -f\n\n")

	vars := make([]*expandString, 0, len(x.expanded))

	for name, v := range x.expanded {
		inserted := false

		// Insert into vars based on dependencies
		for i, vv := range vars {
			if vv.dependsOn(name) {
				tail := make([]*expandString, len(vars)-i)
				copy(tail, vars[i:])

				if i == 0 {
					vars = append([]*expandString{v}, vars...)
				} else {
					vars = append(append(vars[0:i], v), tail...)
				}

				inserted = true
				break
			}
		}

		if !inserted {
			vars = append(vars, v)
		}
	}

	io.WriteString(writer, "# Variables\n")

	for _, v := range vars {
		fmt.Fprintf(writer, "%s ?= ", v.Name)

		for _, part := range v.Parts {
			if part.IsVariable {
				fmt.Fprintf(writer, "$(%s)", part.Value)
			} else {
				fmt.Fprintf(writer, "%s", part.Value)
			}
		}

		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "version ?= ")

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ".")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	io.WriteString(writer, "\n")
	fmt.Fprintf(writer, "major_version = %v\n", Version[0])

	if len(Version) > 1 {
		fmt.Fprintf(writer, "minor_version = %v\n", Version[1])
	}

	if len(Version) > 2 {
		fmt.Fprintf(writer, "micro_version = %v\n", Version[2])
	}

	io.WriteString(writer, "\n")

	target := Target

	if len(target) == 0 {
		pc := make([]uintptr, 3)
		n := runtime.Callers(1, pc)

		me, _ := runtime.FuncForPC(pc[0]).FileLine(pc[0])

		for i := 1; i < n; i++ {
			f := runtime.FuncForPC(pc[i])
			fname, _ := f.FileLine(pc[i])

			if fname != me {
				target = path.Base(path.Dir(fname))
				break
			}
		}
	}

	fmt.Fprintf(writer, "TARGET ?= %s\n", target)

	io.WriteString(writer, "\nSOURCES ?=")
	io.WriteString(writer, "\nSOURCES += $(wildcard *.go)")
	io.WriteString(writer, "\nSOURCES_UNIQUE = $(sort $(SOURCES))")

	io.WriteString(writer, "\n\n")

	io.WriteString(writer, "# Rules\n")
	io.WriteString(writer, "$(TARGET): $(SOURCES_UNIQUE)\n")
	io.WriteString(writer, "\tgo build -o $@\n\n")

	io.WriteString(writer, "clean:\n")
	io.WriteString(writer, "\trm -f $(TARGET)\n\n")

	io.WriteString(writer, "distclean: clean\n\n")

	io.WriteString(writer, "$(TARGET)_installdir ?= $(bindir)\n\n")

	io.WriteString(writer, "install: $(TARGET)\n")
	io.WriteString(writer, "\tmkdir -p $(DESTDIR)$($(TARGET)_installdir) && cp $(TARGET) $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, "uninstall:\n")
	io.WriteString(writer, "\trm -f $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, ".PHONY: install uninstall distclean clean")
}
<|endoftext|>


<commit_before>package fnet

import (
	"errors"
	"sync"
)

const recvBufSize = 20 // This is a total guess as to a reasonable buffer size for our recveive channel

type RoundRobin struct {
	pool      []*FrameConn
	poolLock  sync.Mutex // Lock for changing the pool and pool related values (numConn)
	frameSize int
	numConn   int
	nextConn  int         // The index of the next stream to send a packet on
	recv      chan []byte // Queue of received packets
	wg        sync.WaitGroup
	stopCh    chan struct{}
}

func (rr *RoundRobin) AddConn(fc *FrameConn) {
	rr.poolLock.Lock()
	rr.numConn++
	rr.pool = append(rr.pool, fc)
	rr.poolLock.Unlock()
	// Start a new thread for listening to every connection
	rr.wg.Add(1)
	go rr.listen(fc)
}

func NewRoundRobin(frameSize int) *RoundRobin {
	var conn []*FrameConn
	var wg sync.WaitGroup
	var lock sync.Mutex
	rr := &RoundRobin{
		pool:      conn,
		poolLock:  lock,
		frameSize: frameSize,
		numConn:   0,
		nextConn:  0,
		recv:      make(chan []byte, recvBufSize),
		wg:        wg,
		stopCh:    make(chan struct{}),
	}
	return rr
}

// FrameSize implements FrameConn.FrameSize
func (rr *RoundRobin) FrameSize() int {
	return rr.frameSize
}

// TODO: Make it so that this can be called more than once freely
func (rr *RoundRobin) Stop() {
	close(rr.stopCh)
	for _, conn := range rr.pool {
		(*conn).Stop()
	}
	rr.wg.Wait()
}

// listen for incoming packets and add them to the received queue
func (rr *RoundRobin) listen(fc *FrameConn) {
	defer rr.wg.Done()
	for {
		buf := make([]byte, rr.frameSize)
		sz, err := (*fc).RecvFrame(buf)
		if err != nil {
			select {
			case <-rr.stopCh: // Stop this thread
				return
			default:
				// Remove the stream if the connection is sad
				rr.RemoveConn(fc)
				return
			}
		}
		rr.recv <- buf[:sz]
	}
}

// TODO: Implement this more efficiently
func (rr *RoundRobin) RemoveConn(fc *FrameConn) {
	(*fc).Stop()
	rr.poolLock.Lock()

	// Get index of stream
	// TODO: Exception if doesn't exist at all
	var fcIndex int
	for index, conn := range rr.pool {
		if conn == fc {
			fcIndex = index
			break
		}
	}

	rr.numConn--
	if rr.nextConn >= rr.numConn {
		rr.nextConn = 0
	}
	rr.pool = append(rr.pool[:fcIndex], rr.pool[fcIndex+1:]...)
	rr.poolLock.Unlock()
}

// SendFrame implements FrameConn.SendFrame
func (rr *RoundRobin) SendFrame(b []byte) error {
	rr.poolLock.Lock()
	if rr.numConn == 0 {
		return errors.New("No streams to send packets on.")
	}
	fc := rr.pool[rr.nextConn]
	err := (*fc).SendFrame(b)
	if err != nil {
		rr.poolLock.Unlock()
		rr.RemoveConn(fc)
		if rr.numConn == 0 {
			return errors.New("No streams to send packets on.")
		}
	}
	rr.nextConn = (rr.nextConn + 1) % rr.numConn // Get the next round-robin index
	rr.poolLock.Unlock()
	return nil
}

// RecvFrame implements FrameConn.RecvFrame
// It pulls the next frame out of the recv channel
// This method should be running continously to prevent blocking on the recv chan
func (rr *RoundRobin) RecvFrame(b []byte) (int, error) {
	for {
		select {
		case <-rr.stopCh: // Stop this thread
			return 0, errors.New("Stream stopped.")
		case frame := <-rr.recv:
			copy(b[:len(frame)], frame)

			return len(frame), nil
		}
	}
}
<commit_msg>Actually fixed deadlock problem<commit_after>package fnet

import (
	"errors"
	"sync"
)

const recvBufSize = 20 // This is a total guess as to a reasonable buffer size for our recveive channel

type RoundRobin struct {
	pool      []*FrameConn
	poolLock  sync.Mutex // Lock for changing the pool and pool related values (numConn)
	frameSize int
	numConn   int
	nextConn  int         // The index of the next stream to send a packet on
	recv      chan []byte // Queue of received packets
	wg        sync.WaitGroup
	stopCh    chan struct{}
}

func (rr *RoundRobin) AddConn(fc *FrameConn) {
	rr.poolLock.Lock()
	rr.numConn++
	rr.pool = append(rr.pool, fc)
	rr.poolLock.Unlock()
	// Start a new thread for listening to every connection
	rr.wg.Add(1)
	go rr.listen(fc)
}

func NewRoundRobin(frameSize int) *RoundRobin {
	var conn []*FrameConn
	var wg sync.WaitGroup
	var lock sync.Mutex
	rr := &RoundRobin{
		pool:      conn,
		poolLock:  lock,
		frameSize: frameSize,
		numConn:   0,
		nextConn:  0,
		recv:      make(chan []byte, recvBufSize),
		wg:        wg,
		stopCh:    make(chan struct{}),
	}
	return rr
}

// FrameSize implements FrameConn.FrameSize
func (rr *RoundRobin) FrameSize() int {
	return rr.frameSize
}

// TODO: Make it so that this can be called more than once freely
func (rr *RoundRobin) Stop() {
	close(rr.stopCh)
	for _, conn := range rr.pool {
		(*conn).Stop()
	}
	rr.wg.Wait()
}

// listen for incoming packets and add them to the received queue
func (rr *RoundRobin) listen(fc *FrameConn) {
	defer rr.wg.Done()
	for {
		buf := make([]byte, rr.frameSize)
		sz, err := (*fc).RecvFrame(buf)
		if err != nil {
			select {
			case <-rr.stopCh: // Stop this thread
				return
			default:
				// Remove the stream if the connection is sad
				rr.RemoveConn(fc)
				return
			}
		}
		rr.recv <- buf[:sz]
	}
}

// TODO: Implement this more efficiently
func (rr *RoundRobin) RemoveConn(fc *FrameConn) {
	(*fc).Stop()
	rr.poolLock.Lock()

	// Get index of stream
	// TODO: Exception if doesn't exist at all
	var fcIndex int
	for index, conn := range rr.pool {
		if conn == fc {
			fcIndex = index
			break
		}
	}

	rr.numConn--
	if rr.nextConn >= rr.numConn {
		rr.nextConn = 0
	}
	rr.pool = append(rr.pool[:fcIndex], rr.pool[fcIndex+1:]...)
	rr.poolLock.Unlock()
}

// SendFrame implements FrameConn.SendFrame
func (rr *RoundRobin) SendFrame(b []byte) error {
	rr.poolLock.Lock()
	if rr.numConn == 0 {
		return errors.New("No streams to send packets on.")
	}
	fc := rr.pool[rr.nextConn]
	err := (*fc).SendFrame(b)
	if err != nil {
		rr.poolLock.Unlock()
		rr.RemoveConn(fc)
		return err
	} else {
		rr.nextConn = (rr.nextConn + 1) % rr.numConn // Get the next round-robin index
		rr.poolLock.Unlock()
		return nil
	}
}

// RecvFrame implements FrameConn.RecvFrame
// It pulls the next frame out of the recv channel
// This method should be running continously to prevent blocking on the recv chan
func (rr *RoundRobin) RecvFrame(b []byte) (int, error) {
	for {
		select {
		case <-rr.stopCh: // Stop this thread
			return 0, errors.New("Stream stopped.")
		case frame := <-rr.recv:
			copy(b[:len(frame)], frame)

			return len(frame), nil
		}
	}
}
<|endoftext|>


<commit_before>package simpleforce

import (
	"fmt"
	"os"
	"testing"
)

var (
	force Force
)

func init() {
	session := os.Getenv("FORCE_SESSION")
	url := os.Getenv("FORCE_URL")
	force = New(session, url)
}

type Account struct {
	Name string
}

type Contact struct {
	FirstName string
	LastName  string
	Account   *Account
}

func Example() {
	type Account struct {
		Name string
	}

	type Contact struct {
		FirstName string
		LastName  string
		Account   *Account
	}

	var cs []Contact
	q := force.NewQuery(&cs)
	q.AddConstraint(NewConstraint("Name").EqualsString("Jake Basile"))
	q.Run()
	for _, c := range cs {
		fmt.Printf("%v %v Is From %v", c.FirstName, c.LastName, c.Account.Name)
	}
	// Output:
	// Jake Basile Is From Mutual Mobile
}

func ExampleForce_RunRawQuery() {
	type Contact struct {
		Name string
	}

	var cs []Contact
	force.RunRawQuery("SELECT Name FROM Contact WHERE FirstName='Jake' AND LastName='Basile'", &cs)
	for _, c := range cs {
		fmt.Println(c.Name)
	}
	// Output:
	// Jake Basile
}

func ExampleConstraint() {
	c1 := NewConstraint("FirstName").EqualsString("Jake")
	c2 := NewConstraint("LastName").NotEqualsString("Basile")
	c3 := NewConstraint(c1).Or(c2)
	fmt.Println(c1.Collapse())
	fmt.Println(c2.Collapse())
	fmt.Println(c3.Collapse())
	// Output:
	// (FirstName='Jake')
	// (LastName<>'Basile')
	// ((FirstName='Jake') OR (LastName<>'Basile'))
}

func TestQueryCreation(t *testing.T) {
	var as []Account
	q := force.NewQuery(&as)
	t.Log(q)
}

func TestSimpleConstraintCreation(t *testing.T) {
	c := NewConstraint("FirstName").EqualsString("Jake")
	t.Log(c)
	t.Log(c.Collapse())
	if c.Collapse() != "(FirstName='Jake')" {
		t.Fail()
	}
}

func TestComplexConstraintCreation(t *testing.T) {
	c1 := NewConstraint("FirstName").EqualsString("Jake")
	c2 := NewConstraint("LastName").NotEqualsString("Basile")
	ca := NewConstraint(c1).And(c2)
	c3 := NewConstraint("Account.Name").EqualsString("Mutual Mobile")
	co := NewConstraint(ca).Or(c3)
	t.Log(co)
	t.Log(co.Collapse())
	if co.Collapse() != "(((FirstName='Jake') AND (LastName<>'Basile')) OR (Account.Name='Mutual Mobile'))" {
		t.Fail()
	}
}

func TestSimpleQueryGeneration(t *testing.T) {
	var cs []Contact
	q := force.NewQuery(&cs)
	q.AddConstraint(NewConstraint("FirstName").EqualsString("Jake"))
	q.AddConstraint(NewConstraint("LastName").EqualsString("Basile"))
	q.AddConstraint(NewConstraint("Account.Name").EqualsString("Mutual Mobile"))
	t.Log(q.Generate())
	if q.Generate() != "SELECT FirstName,LastName,Account.Name FROM Contact WHERE (FirstName='Jake') AND (LastName='Basile') AND (Account.Name='Mutual Mobile') LIMIT 10" {
		t.Fail()
	}
}

func TestSimpleQueryRun(t *testing.T) {
	var cs []Contact
	q := force.NewQuery(&cs)
	q.AddConstraint(NewConstraint("Account.Name").NotEqualsString(""))
	q.AddConstraint(NewConstraint("FirstName").NotEqualsString(""))
	q.AddConstraint(NewConstraint("LastName").NotEqualsString(""))
	q.Limit(1000)
	t.Log(q.Generate())
	q.Run()
	for _, c := range cs {
		t.Log(c.FirstName, c.LastName, c.Account.Name)
		if c.FirstName == "" || c.LastName == "" || c.Account == nil || c.Account.Name == "" {
			t.Fail()
		}
	}
}

func TestRawQuery(t *testing.T) {
	var cs []Contact
	force.RunRawQuery("SELECT FirstName FROM Contact WHERE FirstName<>'' LIMIT 1", &cs)
	for _, c := range cs {
		t.Log(c)
		if c.FirstName == "" || c.LastName != "" || c.Account == nil || c.Account.Name != "" {
			t.Fail()
		}
	}
}
<commit_msg>added a benchmark for fun.<commit_after>package simpleforce

import (
	"fmt"
	"os"
	"testing"
)

var (
	force Force
)

func init() {
	session := os.Getenv("FORCE_SESSION")
	url := os.Getenv("FORCE_URL")
	force = New(session, url)
}

type Account struct {
	Name string
}

type Contact struct {
	FirstName string
	LastName  string
	Account   *Account
}

func BenchmarkQuery(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var cs []Contact
		q := force.NewQuery(&cs)
		q.AddConstraint(NewConstraint("Account.Name").NotEqualsString(""))
		q.AddConstraint(NewConstraint("FirstName").NotEqualsString(""))
		q.AddConstraint(NewConstraint("LastName").NotEqualsString(""))
		q.Limit(1000)
		q.Run()
	}
}

func Example() {
	type Account struct {
		Name string
	}

	type Contact struct {
		FirstName string
		LastName  string
		Account   *Account
	}

	var cs []Contact
	q := force.NewQuery(&cs)
	q.AddConstraint(NewConstraint("Name").EqualsString("Jake Basile"))
	q.Run()
	for _, c := range cs {
		fmt.Printf("%v %v Is From %v", c.FirstName, c.LastName, c.Account.Name)
	}
	// Output:
	// Jake Basile Is From Mutual Mobile
}

func ExampleForce_RunRawQuery() {
	type Contact struct {
		Name string
	}

	var cs []Contact
	force.RunRawQuery("SELECT Name FROM Contact WHERE FirstName='Jake' AND LastName='Basile'", &cs)
	for _, c := range cs {
		fmt.Println(c.Name)
	}
	// Output:
	// Jake Basile
}

func ExampleConstraint() {
	c1 := NewConstraint("FirstName").EqualsString("Jake")
	c2 := NewConstraint("LastName").NotEqualsString("Basile")
	c3 := NewConstraint(c1).Or(c2)
	fmt.Println(c1.Collapse())
	fmt.Println(c2.Collapse())
	fmt.Println(c3.Collapse())
	// Output:
	// (FirstName='Jake')
	// (LastName<>'Basile')
	// ((FirstName='Jake') OR (LastName<>'Basile'))
}

func TestQueryCreation(t *testing.T) {
	var as []Account
	q := force.NewQuery(&as)
	t.Log(q)
}

func TestSimpleConstraintCreation(t *testing.T) {
	c := NewConstraint("FirstName").EqualsString("Jake")
	t.Log(c)
	t.Log(c.Collapse())
	if c.Collapse() != "(FirstName='Jake')" {
		t.Fail()
	}
}

func TestComplexConstraintCreation(t *testing.T) {
	c1 := NewConstraint("FirstName").EqualsString("Jake")
	c2 := NewConstraint("LastName").NotEqualsString("Basile")
	ca := NewConstraint(c1).And(c2)
	c3 := NewConstraint("Account.Name").EqualsString("Mutual Mobile")
	co := NewConstraint(ca).Or(c3)
	t.Log(co)
	t.Log(co.Collapse())
	if co.Collapse() != "(((FirstName='Jake') AND (LastName<>'Basile')) OR (Account.Name='Mutual Mobile'))" {
		t.Fail()
	}
}

func TestSimpleQueryGeneration(t *testing.T) {
	var cs []Contact
	q := force.NewQuery(&cs)
	q.AddConstraint(NewConstraint("FirstName").EqualsString("Jake"))
	q.AddConstraint(NewConstraint("LastName").EqualsString("Basile"))
	q.AddConstraint(NewConstraint("Account.Name").EqualsString("Mutual Mobile"))
	t.Log(q.Generate())
	if q.Generate() != "SELECT FirstName,LastName,Account.Name FROM Contact WHERE (FirstName='Jake') AND (LastName='Basile') AND (Account.Name='Mutual Mobile') LIMIT 10" {
		t.Fail()
	}
}

func TestSimpleQueryRun(t *testing.T) {
	var cs []Contact
	q := force.NewQuery(&cs)
	q.AddConstraint(NewConstraint("Account.Name").NotEqualsString(""))
	q.AddConstraint(NewConstraint("FirstName").NotEqualsString(""))
	q.AddConstraint(NewConstraint("LastName").NotEqualsString(""))
	q.Limit(1000)
	t.Log(q.Generate())
	q.Run()
	if len(cs) != 1000 {
		t.Fail()
	}
	for _, c := range cs {
		t.Log(c.FirstName, c.LastName, c.Account.Name)
		if c.FirstName == "" || c.LastName == "" || c.Account == nil || c.Account.Name == "" {
			t.Fail()
		}
	}
}

func TestRawQuery(t *testing.T) {
	var cs []Contact
	force.RunRawQuery("SELECT FirstName FROM Contact WHERE FirstName<>'' LIMIT 1", &cs)
	for _, c := range cs {
		t.Log(c)
		if c.FirstName == "" || c.LastName != "" || c.Account == nil || c.Account.Name != "" {
			t.Fail()
		}
	}
}
<|endoftext|>


<commit_before>package worker

import (
	"github.com/Sirupsen/logrus"
	"github.com/confluentinc/confluent-kafka-go/kafka"
	"sync"
	"time"
)

// MessageProcessor will ne called when new message arrived
type MessageProcessor interface {
	ProcessMessage(*kafka.Message) error
	StartDaemon() error
	StopDaemon() error
}

type ConsumerWorker struct {
	broker            string
	group             string
	topic             string
	threads           int64
	logger            *logrus.Entry
	stopChan          chan bool
	waitGroup         *sync.WaitGroup
	statsMutex        *sync.Mutex
	errors            int64
	succeed           int64
	messageProcessor  MessageProcessor // Called after new message arrived
	workerThreads     []*workerThread
	reportStatsTicker *time.Ticker
}

func (w *ConsumerWorker) Stop() error {

	w.logger.Infof("begin stop consumer worker daemonThread")
	if err := w.messageProcessor.StopDaemon(); err != nil {
		w.logger.Errorf("failt to StopDaemon %s", err)
	}

	w.logger.Infof("begin stop consumer worker thread")
	for _, wt := range w.workerThreads {
		wt.stop()
	}
	w.logger.Infof("end stop consumer worker thread")

	close(w.stopChan)

	return nil
}

func (w *ConsumerWorker) Wait() {
	w.waitGroup.Wait()
}

func (w *ConsumerWorker) Start() error {
	if err := w.messageProcessor.StartDaemon(); err != nil {
		return err
	}

	var i int64
	for i = 0; i < w.threads; i++ {
		wt := newWorkerThread(w, i)
		logger := w.logger.WithField("worker-id", wt.id)
		wt.logger = logger
		w.workerThreads = append(w.workerThreads, wt)
		w.waitGroup.Add(1)
		go wt.start()
	}

	go func() {
	outer:
		for {
			select {
			case <-w.reportStatsTicker.C:
				w.printStats()
			case <-w.stopChan:
				break outer
			}
		}
		w.logger.Infof("ConsumerWorker report stats goroutine exit")
	}()

	return nil
}

func NewConsumerWorker(broker, group, topic string,
	threads int64, messageProcessor MessageProcessor,
	logger *logrus.Entry) (*ConsumerWorker, error) {

	w := ConsumerWorker{
		broker:            broker,
		group:             group,
		topic:             topic,
		threads:           threads,
		messageProcessor:  messageProcessor,
		waitGroup:         &sync.WaitGroup{},
		logger:            logger,
		workerThreads:     []*workerThread{},
		errors:            0,
		succeed:           0,
		statsMutex:        &sync.Mutex{},
		reportStatsTicker: time.NewTicker(60 * time.Second),
	}
	w.stopChan = make(chan bool)
	logger.Debugf("NewConsumerWorker(broker: %s, group: %s, topic: %s, threads: %d)",
		broker, group, topic, threads)
	return &w, nil
}

func (w *ConsumerWorker) printStats() {
	w.statsMutex.Lock()
	defer w.statsMutex.Unlock()
	w.logger.Infof("processed %d, error %d, succeed %d", (w.errors + w.succeed), w.errors, w.succeed)
	w.errors = 0
	w.succeed = 0
}

func (w *ConsumerWorker) IncErrors(errors int) {
	w.statsMutex.Lock()
	defer w.statsMutex.Unlock()
	w.errors = w.errors + int64(errors)
}

func (w *ConsumerWorker) IncSucceed(succeed int) {
	w.statsMutex.Lock()
	defer w.statsMutex.Unlock()
	w.succeed = w.succeed + int64(succeed)

}
<commit_msg>fix: avoid app not exist<commit_after>package worker

import (
	"github.com/Sirupsen/logrus"
	"github.com/confluentinc/confluent-kafka-go/kafka"
	"sync"
	"time"
)

// MessageProcessor will ne called when new message arrived
type MessageProcessor interface {
	ProcessMessage(*kafka.Message) error
	StartDaemon() error
	StopDaemon() error
}

type ConsumerWorker struct {
	broker            string
	group             string
	topic             string
	threads           int64
	logger            *logrus.Entry
	stopChan          chan bool
	waitGroup         *sync.WaitGroup
	statsMutex        *sync.Mutex
	errors            int64
	succeed           int64
	messageProcessor  MessageProcessor // Called after new message arrived
	workerThreads     []*workerThread
	reportStatsTicker *time.Ticker
}

func (w *ConsumerWorker) Stop() error {

	w.logger.Infof("begin stop consumer worker daemonThread")
	if err := w.messageProcessor.StopDaemon(); err != nil {
		w.logger.Errorf("failt to StopDaemon %s", err)
	}

	w.logger.Infof("begin stop consumer worker thread")
	for _, wt := range w.workerThreads {
		w.waitGroup.Done()
		wt.stop()
	}
	w.logger.Infof("end stop consumer worker thread")

	close(w.stopChan)

	return nil
}

func (w *ConsumerWorker) Wait() {
	w.waitGroup.Wait()
}

func (w *ConsumerWorker) Start() error {
	if err := w.messageProcessor.StartDaemon(); err != nil {
		return err
	}

	var i int64
	for i = 0; i < w.threads; i++ {
		wt := newWorkerThread(w, i)
		logger := w.logger.WithField("worker-id", wt.id)
		wt.logger = logger
		w.workerThreads = append(w.workerThreads, wt)
		w.waitGroup.Add(1)
		go wt.start()
	}

	go func() {
	outer:
		for {
			select {
			case <-w.reportStatsTicker.C:
				w.printStats()
			case <-w.stopChan:
				break outer
			}
		}
		w.logger.Infof("ConsumerWorker report stats goroutine exit")
	}()

	return nil
}

func NewConsumerWorker(broker, group, topic string,
	threads int64, messageProcessor MessageProcessor,
	logger *logrus.Entry) (*ConsumerWorker, error) {

	w := ConsumerWorker{
		broker:            broker,
		group:             group,
		topic:             topic,
		threads:           threads,
		messageProcessor:  messageProcessor,
		waitGroup:         &sync.WaitGroup{},
		logger:            logger,
		workerThreads:     []*workerThread{},
		errors:            0,
		succeed:           0,
		statsMutex:        &sync.Mutex{},
		reportStatsTicker: time.NewTicker(60 * time.Second),
	}
	w.stopChan = make(chan bool)
	logger.Debugf("NewConsumerWorker(broker: %s, group: %s, topic: %s, threads: %d)",
		broker, group, topic, threads)
	return &w, nil
}

func (w *ConsumerWorker) printStats() {
	w.statsMutex.Lock()
	defer w.statsMutex.Unlock()
	w.logger.Infof("processed %d, error %d, succeed %d", (w.errors + w.succeed), w.errors, w.succeed)
	w.errors = 0
	w.succeed = 0
}

func (w *ConsumerWorker) IncErrors(errors int) {
	w.statsMutex.Lock()
	defer w.statsMutex.Unlock()
	w.errors = w.errors + int64(errors)
}

func (w *ConsumerWorker) IncSucceed(succeed int) {
	w.statsMutex.Lock()
	defer w.statsMutex.Unlock()
	w.succeed = w.succeed + int64(succeed)

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"net/url"

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

type natsMetrics struct {
	Connections float64 `json:"connections"`
	Routes      float64 `json:"routes"`

	MessagesIn  float64 `json:"in_msgs"`
	MessagesOut float64 `json:"out_msgs"`

	BytesIn  float64 `json:"in_bytes"`
	BytesOut float64 `json:"out_bytes"`

	SlowConsumers float64 `json:"slow_consumers"`
}

var (
	connections = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "connections",
		Help:      "Active connections to gnatsd",
	})

	routes = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "routes",
		Help:      "Active routes to gnatsd",
	})

	messageCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "msg_total",
		Help:      "Count of transferred messages",
	},
		[]string{"direction"},
	)

	bytesCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "bytes_total",
		Help:      "Amount of transmitted data",
	},
		[]string{"direction"},
	)

	slowConsumers = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "slow_consumers",
		Help:      "Number of slow consumers",
	})
)

type Exporter struct {
	NatsURL string
}

func NewExporter(natsURL *url.URL) *Exporter {
	natsURL.Path = "/varz"
	return &Exporter{
		NatsURL: natsURL.String(),
	}
}

func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	connections.Describe(ch)
	routes.Describe(ch)
	messageCounter.Describe(ch)
	bytesCounter.Describe(ch)
	slowConsumers.Describe(ch)
}

func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.collect()

	connections.Collect(ch)
	routes.Collect(ch)
	messageCounter.Collect(ch)
	bytesCounter.Collect(ch)
	slowConsumers.Collect(ch)
}

func (e *Exporter) collect() {
	var metrics natsMetrics
	resp, err := http.DefaultClient.Get(e.NatsURL)
	if err != nil {
		log.Printf("could not retrieve NATS metrics: %v", err)
		return
	}

	err = json.NewDecoder(resp.Body).Decode(&metrics)
	if err != nil {
		log.Printf("could not decode NATS metrics: %v", err)
	}

	connections.Set(metrics.Connections)
	routes.Set(metrics.Routes)

	messageCounter.WithLabelValues("in").Add(metrics.MessagesIn)
	messageCounter.WithLabelValues("out").Add(metrics.MessagesOut)

	bytesCounter.WithLabelValues("in").Add(metrics.BytesIn)
	bytesCounter.WithLabelValues("out").Add(metrics.BytesOut)

	slowConsumers.Set(metrics.SlowConsumers)
}
<commit_msg>added method documentation<commit_after>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"net/url"
	"time"

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

type natsMetrics struct {
	Connections float64 `json:"connections"`
	Routes      float64 `json:"routes"`

	MessagesIn  float64 `json:"in_msgs"`
	MessagesOut float64 `json:"out_msgs"`

	BytesIn  float64 `json:"in_bytes"`
	BytesOut float64 `json:"out_bytes"`

	SlowConsumers float64 `json:"slow_consumers"`
}

var (
	connections = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "connections",
		Help:      "Active connections to gnatsd",
	})

	routes = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "routes",
		Help:      "Active routes to gnatsd",
	})

	messageCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "msg_total",
		Help:      "Count of transferred messages",
	},
		[]string{"direction"},
	)

	bytesCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "bytes_total",
		Help:      "Amount of transmitted data",
	},
		[]string{"direction"},
	)

	slowConsumers = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "slow_consumers",
		Help:      "Number of slow consumers",
	})
)

// Exporter implements the prometheus.Collector interface. It exposes the metrics
// of a NATS node.
type Exporter struct {
	NatsURL string
}

// NewExporter instantiates a new NATS Exporter.
func NewExporter(natsURL *url.URL) *Exporter {
	natsURL.Path = "/varz"
	return &Exporter{
		NatsURL: natsURL.String(),
	}
}

// Describe describes all the registered stats metrics from the NATS node.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	connections.Describe(ch)
	routes.Describe(ch)
	messageCounter.Describe(ch)
	bytesCounter.Describe(ch)
	slowConsumers.Describe(ch)
}

// Collect collects all the registered stats metrics from the NATS node.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.collect()

	connections.Collect(ch)
	routes.Collect(ch)
	messageCounter.Collect(ch)
	bytesCounter.Collect(ch)
	slowConsumers.Collect(ch)
}

func (e *Exporter) collect() {
	var metrics natsMetrics

	httpClient := http.DefaultClient
	httpClient.Timeout = 1 * time.Second
	resp, err := httpClient.Get(e.NatsURL)
	if err != nil {
		log.Printf("could not retrieve NATS metrics: %v", err)
		return
	}

	err = json.NewDecoder(resp.Body).Decode(&metrics)
	if err != nil {
		log.Printf("could not decode NATS metrics: %v", err)
	}

	connections.Set(metrics.Connections)
	routes.Set(metrics.Routes)

	messageCounter.WithLabelValues("in").Add(metrics.MessagesIn)
	messageCounter.WithLabelValues("out").Add(metrics.MessagesOut)

	bytesCounter.WithLabelValues("in").Add(metrics.BytesIn)
	bytesCounter.WithLabelValues("out").Add(metrics.BytesOut)

	slowConsumers.Set(metrics.SlowConsumers)
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 librdkafka

// #cgo CFLAGS: -I/usr/local/include -std=c99
// #cgo LDFLAGS: -L/usr/local/opt/librdkafka/lib -L/usr/local/lib -lrdkafka
// #include "wrapper.h"
import "C"

// Topic wrapper handle for rd_kafka_topic_t
type Topic struct {
	handle   *C.rd_kafka_topic_t
	client   *Client
	name     string
	shutdown bool
}

// NewTopic creates a new topic representation in librdkafka.
// You have to call Close() to free any internal state. As this struct holds a
// pointer to the client make sure that Client.Close is called after closing
// objects of this type.
func NewTopic(name string, config TopicConfig, client *Client) *Topic {
	return &Topic{
		handle:   C.rd_kafka_topic_new(client.handle, C.CString(name), config.handle),
		client:   client,
		name:     name,
		shutdown: false,
	}
}

// TriggerShutdown signals a topic to stop producing messages (unblocks any
// waiting topics).
func (t *Topic) TriggerShutdown() {
	t.shutdown = true
}

// Close frees the internal handle and tries to flush the queue.
func (t *Topic) Close() {
	oldQueueLen := C.int(0x7FFFFFFF)
	queueLen := C.rd_kafka_outq_len(t.client.handle)

	// Wait as long as we're flushing
	for queueLen > 0 && queueLen < oldQueueLen {
		C.rd_kafka_poll(t.client.handle, 1000)
		oldQueueLen = queueLen
		queueLen = C.rd_kafka_outq_len(t.client.handle)
	}

	if queueLen > 0 {
		Log.Printf("%d messages have been lost as the internal queue could not be flushed", queueLen)
	}
	C.rd_kafka_topic_destroy(t.handle)
}

// GetName returns the name of the topic
func (t *Topic) GetName() string {
	return t.name
}

// Poll polls for new data to be sent to the async handler functions
func (t *Topic) Poll() {
	C.rd_kafka_poll(t.client.handle, 1000)
}

// Produce produces a single messages.
// If a message cannot be produced because of internal (non-wire) problems an
// error is immediately returned instead of being asynchronously handled via
// MessageDelivery interface.
func (t *Topic) Produce(message Message) error {
	keyLen, keyPtr, payLen, payPtr, usrLen, usrPtr := MarshalMessage(message)
	usrData := C.CreateBuffer(usrLen, usrPtr)
	success := C.rd_kafka_produce(t.handle, C.RD_KAFKA_PARTITION_UA, C.RD_KAFKA_MSG_F_COPY, payPtr, payLen, keyPtr, keyLen, usrData)

	if success != 0 {
		defer C.DestroyBuffer(usrData)
		rspErr := ResponseError{
			Userdata: message.GetUserdata(),
			Code:     int(C.GetLastError()),
		}
		return rspErr // ### return, error ###
	}

	C.rd_kafka_poll(t.client.handle, 0)
	return nil
}

// ProduceBatch produces a set of messages.
// Messages that cannot be produced because of internal (non-wire) problems are
// immediately returned instead of asynchronously handled via MessageDelivery
// interface.
func (t *Topic) ProduceBatch(messages []Message) []error {
	errors := []error{}
	if len(messages) == 0 {
		return errors // ### return, nothing to do ###
	}

	batch := PrepareBatch(messages)
	batchLen := C.int(len(messages))
	defer C.DestroyBatch(batch)

	enqueued := C.rd_kafka_produce_batch(t.handle, C.RD_KAFKA_PARTITION_UA, C.RD_KAFKA_MSG_F_COPY, batch, batchLen)
	if enqueued != batchLen {
		offset := C.int(0)
		for offset >= 0 {
			offset = C.BatchGetNextError(batch, batchLen, offset)
			if offset >= 0 {
				bufferPtr := C.BatchGetUserdataAt(batch, offset)
				errCode := C.BatchGetErrAt(batch, offset)

				rspErr := ResponseError{
					Userdata: UnmarshalBuffer(bufferPtr),
					Code:     int(errCode),
				}

				errors = append(errors, rspErr)
				offset++
			}
		}
	}

	for C.rd_kafka_outq_len(t.client.handle) > 0 && !t.shutdown {
		C.rd_kafka_poll(t.client.handle, 20)
	}

	return errors
}
<commit_msg>cgo 1.5 fixes<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 librdkafka

// #cgo CFLAGS: -I/usr/local/include -std=c99
// #cgo LDFLAGS: -L/usr/local/opt/librdkafka/lib -L/usr/local/lib -lrdkafka
// #include "wrapper.h"
import "C"

// Topic wrapper handle for rd_kafka_topic_t
type Topic struct {
	handle   *C.rd_kafka_topic_t
	client   *Client
	name     string
	shutdown bool
}

// NewTopic creates a new topic representation in librdkafka.
// You have to call Close() to free any internal state. As this struct holds a
// pointer to the client make sure that Client.Close is called after closing
// objects of this type.
func NewTopic(name string, config TopicConfig, client *Client) *Topic {
	return &Topic{
		handle:   C.rd_kafka_topic_new(client.handle, C.CString(name), config.handle),
		client:   client,
		name:     name,
		shutdown: false,
	}
}

// TriggerShutdown signals a topic to stop producing messages (unblocks any
// waiting topics).
func (t *Topic) TriggerShutdown() {
	t.shutdown = true
}

// Close frees the internal handle and tries to flush the queue.
func (t *Topic) Close() {
	oldQueueLen := C.int(0x7FFFFFFF)
	queueLen := C.rd_kafka_outq_len(t.client.handle)

	// Wait as long as we're flushing
	for queueLen > 0 && queueLen < oldQueueLen {
		C.rd_kafka_poll(t.client.handle, 1000)
		oldQueueLen = queueLen
		queueLen = C.rd_kafka_outq_len(t.client.handle)
	}

	if queueLen > 0 {
		Log.Printf("%d messages have been lost as the internal queue could not be flushed", queueLen)
	}
	C.rd_kafka_topic_destroy(t.handle)
}

// GetName returns the name of the topic
func (t *Topic) GetName() string {
	return t.name
}

// Poll polls for new data to be sent to the async handler functions
func (t *Topic) Poll() {
	C.rd_kafka_poll(t.client.handle, 1000)
}

// Produce produces a single messages.
// If a message cannot be produced because of internal (non-wire) problems an
// error is immediately returned instead of being asynchronously handled via
// MessageDelivery interface.
func (t *Topic) Produce(message Message) error {
	keyLen, keyPtr, payLen, payPtr, usrLen, usrPtr := MarshalMessage(message)
	usrData := C.CreateBuffer(usrLen, usrPtr)
	success := C.rd_kafka_produce(t.handle, C.RD_KAFKA_PARTITION_UA, C.RD_KAFKA_MSG_F_COPY, payPtr, payLen, keyPtr, keyLen, usrData)

	if success != 0 {
		defer C.DestroyBuffer((*C.buffer_t)(usrData))
		rspErr := ResponseError{
			Userdata: message.GetUserdata(),
			Code:     int(C.GetLastError()),
		}
		return rspErr // ### return, error ###
	}

	C.rd_kafka_poll(t.client.handle, 0)
	return nil
}

// ProduceBatch produces a set of messages.
// Messages that cannot be produced because of internal (non-wire) problems are
// immediately returned instead of asynchronously handled via MessageDelivery
// interface.
func (t *Topic) ProduceBatch(messages []Message) []error {
	errors := []error{}
	if len(messages) == 0 {
		return errors // ### return, nothing to do ###
	}

	batch := PrepareBatch(messages)
	batchLen := C.int(len(messages))
	defer C.DestroyBatch(batch)

	enqueued := C.rd_kafka_produce_batch(t.handle, C.RD_KAFKA_PARTITION_UA, C.RD_KAFKA_MSG_F_COPY, batch, batchLen)
	if enqueued != batchLen {
		offset := C.int(0)
		for offset >= 0 {
			offset = C.BatchGetNextError(batch, batchLen, offset)
			if offset >= 0 {
				bufferPtr := C.BatchGetUserdataAt(batch, offset)
				errCode := C.BatchGetErrAt(batch, offset)

				rspErr := ResponseError{
					Userdata: UnmarshalBuffer(bufferPtr),
					Code:     int(errCode),
				}

				errors = append(errors, rspErr)
				offset++
			}
		}
	}

	for C.rd_kafka_outq_len(t.client.handle) > 0 && !t.shutdown {
		C.rd_kafka_poll(t.client.handle, 20)
	}

	return errors
}
<|endoftext|>