<commit_before>package mdata

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"time"

	"github.com/bitly/go-hostpool"
	"github.com/nsqio/go-nsq"
	"github.com/raintank/met"
	clNSQ "github.com/raintank/metrictank/mdata/clnsq"
	"github.com/raintank/misc/instrumented_nsq"
	"github.com/raintank/worldping-api/pkg/log"
)

var (
	hostPool    hostpool.HostPool
	producers   map[string]*nsq.Producer
	Enabled     bool
	nsqdAdds    []string
	lookupdAdds []string
	topic       string
	channel     string
	pCfg        *nsq.Config
	cCfg        *nsq.Config
)

type ClNSQ struct {
	in       chan SavedChunk
	buf      []SavedChunk
	instance string
	Cl
}

func NewNSQ(instance string, metrics Metrics, stats met.Backend) *ClNSQ {
	Enabled = clNSQ.Enabled
	nsqdAdds = clNSQ.NsqdAdds
	lookupdAdds = clNSQ.LookupdAdds
	topic = clNSQ.Topic
	channel = clNSQ.Channel
	pCfg = clNSQ.PCfg
	cCfg = clNSQ.CCfg

	// producers
	hostPool = hostpool.NewEpsilonGreedy(nsqdAdds, 0, &hostpool.LinearEpsilonValueCalculator{})
	producers = make(map[string]*nsq.Producer)

	for _, addr := range nsqdAdds {
		producer, err := nsq.NewProducer(addr, pCfg)
		if err != nil {
			log.Fatal(4, "failed creating producer %s", err.Error())
		}
		producers[addr] = producer
	}

	// consumers
	consumer, err := insq.NewConsumer(topic, channel, cCfg, "metric_persist.%s", stats)
	if err != nil {
		log.Fatal(4, "Failed to create NSQ consumer. %s", err)
	}
	c := &ClNSQ{
		in:       make(chan SavedChunk),
		instance: instance,
		Cl: Cl{
			instance: instance,
			metrics:  metrics,
		},
	}
	consumer.AddConcurrentHandlers(c, 2)

	err = consumer.ConnectToNSQDs(nsqdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQDs. %s", err)
	}
	log.Info("persist consumer connected to nsqd")

	err = consumer.ConnectToNSQLookupds(lookupdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQLookupds. %s", err)
	}
	go c.run()
	return c
}

func (c *ClNSQ) HandleMessage(m *nsq.Message) error {
	c.Handle(m.Body)
	return nil
}

func (c *ClNSQ) Send(sc SavedChunk) {
	c.in <- sc
}

func (c *ClNSQ) run() {
	ticker := time.NewTicker(time.Second)
	max := 5000
	for {
		select {
		case chunk := <-c.in:
			c.buf = append(c.buf, chunk)
			if len(c.buf) == max {
				c.flush()
			}
		case <-ticker.C:
			c.flush()
		}
	}
}

// flush makes sure the batch gets sent, asynchronously.
func (c *ClNSQ) flush() {
	if len(c.buf) == 0 {
		return
	}

	msg := PersistMessageBatch{Instance: c.instance, SavedChunks: c.buf}
	c.buf = nil

	go func() {
		log.Debug("CLU sending %d batch metricPersist messages", len(msg.SavedChunks))

		data, err := json.Marshal(&msg)
		if err != nil {
			log.Fatal(4, "failed to marshal persistMessage to json.")
		}
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.LittleEndian, uint8(PersistMessageBatchV1))
		buf.Write(data)
		messagesSize.Value(int64(buf.Len()))

		sent := false
		for !sent {
			// This will always return a host. If all hosts are currently marked as dead,
			// then all hosts will be reset to alive and we will try them all again. This
			// will result in this loop repeating forever until we successfully publish our msg.
			hostPoolResponse := hostPool.Get()
			prod := producers[hostPoolResponse.Host()]
			err = prod.Publish(topic, buf.Bytes())
			// Hosts that are marked as dead will be retried after 30seconds.  If we published
			// successfully, then sending a nil error will mark the host as alive again.
			hostPoolResponse.Mark(err)
			if err != nil {
				log.Warn("publisher marking host %s as faulty due to %s", hostPoolResponse.Host(), err)
			} else {
				sent = true
			}
			time.Sleep(time.Second)
		}
		messagesPublished.Inc(1)
	}()
}
<commit_msg>this field not needed. we look at clNSQ's package Enabled flag<commit_after>package mdata

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"time"

	"github.com/bitly/go-hostpool"
	"github.com/nsqio/go-nsq"
	"github.com/raintank/met"
	clNSQ "github.com/raintank/metrictank/mdata/clnsq"
	"github.com/raintank/misc/instrumented_nsq"
	"github.com/raintank/worldping-api/pkg/log"
)

var (
	hostPool    hostpool.HostPool
	producers   map[string]*nsq.Producer
	nsqdAdds    []string
	lookupdAdds []string
	topic       string
	channel     string
	pCfg        *nsq.Config
	cCfg        *nsq.Config
)

type ClNSQ struct {
	in       chan SavedChunk
	buf      []SavedChunk
	instance string
	Cl
}

func NewNSQ(instance string, metrics Metrics, stats met.Backend) *ClNSQ {
	nsqdAdds = clNSQ.NsqdAdds
	lookupdAdds = clNSQ.LookupdAdds
	topic = clNSQ.Topic
	channel = clNSQ.Channel
	pCfg = clNSQ.PCfg
	cCfg = clNSQ.CCfg

	// producers
	hostPool = hostpool.NewEpsilonGreedy(nsqdAdds, 0, &hostpool.LinearEpsilonValueCalculator{})
	producers = make(map[string]*nsq.Producer)

	for _, addr := range nsqdAdds {
		producer, err := nsq.NewProducer(addr, pCfg)
		if err != nil {
			log.Fatal(4, "failed creating producer %s", err.Error())
		}
		producers[addr] = producer
	}

	// consumers
	consumer, err := insq.NewConsumer(topic, channel, cCfg, "metric_persist.%s", stats)
	if err != nil {
		log.Fatal(4, "Failed to create NSQ consumer. %s", err)
	}
	c := &ClNSQ{
		in:       make(chan SavedChunk),
		instance: instance,
		Cl: Cl{
			instance: instance,
			metrics:  metrics,
		},
	}
	consumer.AddConcurrentHandlers(c, 2)

	err = consumer.ConnectToNSQDs(nsqdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQDs. %s", err)
	}
	log.Info("persist consumer connected to nsqd")

	err = consumer.ConnectToNSQLookupds(lookupdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQLookupds. %s", err)
	}
	go c.run()
	return c
}

func (c *ClNSQ) HandleMessage(m *nsq.Message) error {
	c.Handle(m.Body)
	return nil
}

func (c *ClNSQ) Send(sc SavedChunk) {
	c.in <- sc
}

func (c *ClNSQ) run() {
	ticker := time.NewTicker(time.Second)
	max := 5000
	for {
		select {
		case chunk := <-c.in:
			c.buf = append(c.buf, chunk)
			if len(c.buf) == max {
				c.flush()
			}
		case <-ticker.C:
			c.flush()
		}
	}
}

// flush makes sure the batch gets sent, asynchronously.
func (c *ClNSQ) flush() {
	if len(c.buf) == 0 {
		return
	}

	msg := PersistMessageBatch{Instance: c.instance, SavedChunks: c.buf}
	c.buf = nil

	go func() {
		log.Debug("CLU sending %d batch metricPersist messages", len(msg.SavedChunks))

		data, err := json.Marshal(&msg)
		if err != nil {
			log.Fatal(4, "failed to marshal persistMessage to json.")
		}
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.LittleEndian, uint8(PersistMessageBatchV1))
		buf.Write(data)
		messagesSize.Value(int64(buf.Len()))

		sent := false
		for !sent {
			// This will always return a host. If all hosts are currently marked as dead,
			// then all hosts will be reset to alive and we will try them all again. This
			// will result in this loop repeating forever until we successfully publish our msg.
			hostPoolResponse := hostPool.Get()
			prod := producers[hostPoolResponse.Host()]
			err = prod.Publish(topic, buf.Bytes())
			// Hosts that are marked as dead will be retried after 30seconds.  If we published
			// successfully, then sending a nil error will mark the host as alive again.
			hostPoolResponse.Mark(err)
			if err != nil {
				log.Warn("publisher marking host %s as faulty due to %s", hostPoolResponse.Host(), err)
			} else {
				sent = true
			}
			time.Sleep(time.Second)
		}
		messagesPublished.Inc(1)
	}()
}
<|endoftext|>


<commit_before>package org

import (
	"regexp"
	"strings"
)

type Keyword struct {
	Key   string
	Value string
}

type NodeWithMeta struct {
	Node Node
	Meta map[string]string
}

type Comment struct{ Content string }

var keywordRegexp = regexp.MustCompile(`^(\s*)#\+([^:]+):\s(.*)`)
var commentRegexp = regexp.MustCompile(`^(\s*)#(.*)`)
var affiliatedKeywordRegexp = regexp.MustCompile(`^(CAPTION)$`)

func lexKeywordOrComment(line string) (token, bool) {
	if m := keywordRegexp.FindStringSubmatch(line); m != nil {
		return token{"keyword", len(m[1]), m[2], m}, true
	} else if m := commentRegexp.FindStringSubmatch(line); m != nil {
		return token{"comment", len(m[1]), m[2], m}, true
	}
	return nilToken, false
}

func (d *Document) parseKeyword(i int, stop stopFn) (int, Node) {
	k := parseKeyword(d.tokens[i])
	if affiliatedKeywordRegexp.MatchString(k.Key) {
		consumed, node := d.parseAffiliated(i, stop)
		if consumed != 0 {
			return consumed, node
		}
	} else {
		d.BufferSettings[k.Key] = strings.Join([]string{d.BufferSettings[k.Key], k.Value}, "\n")
	}
	return 1, k
}

func (d *Document) parseComment(i int, stop stopFn) (int, Node) {
	return 1, Comment{d.tokens[i].content}
}

func (d *Document) parseAffiliated(i int, stop stopFn) (int, Node) {
	start, meta := i, map[string]string{}
	for ; !stop(d, i) && d.tokens[i].kind == "keyword"; i++ {
		k := parseKeyword(d.tokens[i])
		if !affiliatedKeywordRegexp.MatchString(k.Key) {
			return 0, nil
		}
		if value, ok := meta[k.Key]; ok {
			meta[k.Key] = value + " " + k.Value
		} else {
			meta[k.Key] = k.Value
		}
	}
	if stop(d, i) {
		return 0, nil
	}
	consumed, node := 0, (Node)(nil)
	if t := d.tokens[i]; t.kind == "text" {
		if nodes := d.parseInline(t.content); len(nodes) == 1 && isImageOrVideoLink(nodes[0]) {
			consumed, node = 1, Paragraph{nodes[:1]}
		}
	} else {
		consumed, node = d.parseOne(i, stop)
	}
	if consumed == 0 || node == nil {
		return 0, nil
	}
	i += consumed
	return i - start, NodeWithMeta{node, meta}
}

func parseKeyword(t token) Keyword {
	k, v := t.matches[2], t.matches[3]
	k = strings.ToUpper(k)
	return Keyword{k, v}
}
<commit_msg>Support keywords without a value like "#+RESULTS:"<commit_after>package org

import (
	"regexp"
	"strings"
)

type Keyword struct {
	Key   string
	Value string
}

type NodeWithMeta struct {
	Node Node
	Meta map[string]string
}

type Comment struct{ Content string }

var keywordRegexp = regexp.MustCompile(`^(\s*)#\+([^:]+):(\s+(.*)|(\s*)$)`)
var commentRegexp = regexp.MustCompile(`^(\s*)#(.*)`)
var affiliatedKeywordRegexp = regexp.MustCompile(`^(CAPTION)$`)

func lexKeywordOrComment(line string) (token, bool) {
	if m := keywordRegexp.FindStringSubmatch(line); m != nil {
		return token{"keyword", len(m[1]), m[2], m}, true
	} else if m := commentRegexp.FindStringSubmatch(line); m != nil {
		return token{"comment", len(m[1]), m[2], m}, true
	}
	return nilToken, false
}

func (d *Document) parseKeyword(i int, stop stopFn) (int, Node) {
	k := parseKeyword(d.tokens[i])
	if affiliatedKeywordRegexp.MatchString(k.Key) {
		consumed, node := d.parseAffiliated(i, stop)
		if consumed != 0 {
			return consumed, node
		}
	} else {
		d.BufferSettings[k.Key] = strings.Join([]string{d.BufferSettings[k.Key], k.Value}, "\n")
	}
	return 1, k
}

func (d *Document) parseComment(i int, stop stopFn) (int, Node) {
	return 1, Comment{d.tokens[i].content}
}

func (d *Document) parseAffiliated(i int, stop stopFn) (int, Node) {
	start, meta := i, map[string]string{}
	for ; !stop(d, i) && d.tokens[i].kind == "keyword"; i++ {
		k := parseKeyword(d.tokens[i])
		if !affiliatedKeywordRegexp.MatchString(k.Key) {
			return 0, nil
		}
		if value, ok := meta[k.Key]; ok {
			meta[k.Key] = value + " " + k.Value
		} else {
			meta[k.Key] = k.Value
		}
	}
	if stop(d, i) {
		return 0, nil
	}
	consumed, node := 0, (Node)(nil)
	if t := d.tokens[i]; t.kind == "text" {
		if nodes := d.parseInline(t.content); len(nodes) == 1 && isImageOrVideoLink(nodes[0]) {
			consumed, node = 1, Paragraph{nodes[:1]}
		}
	} else {
		consumed, node = d.parseOne(i, stop)
	}
	if consumed == 0 || node == nil {
		return 0, nil
	}
	i += consumed
	return i - start, NodeWithMeta{node, meta}
}

func parseKeyword(t token) Keyword {
	k, v := t.matches[2], t.matches[4]
	k = strings.ToUpper(k)
	return Keyword{k, v}
}
<|endoftext|>


<commit_before>package orgs

import (
	"encoding/json"
	"fmt"
	"github.com/Financial-Times/neo-cypher-runner-go"
	"github.com/jmcvetta/neoism"
	"strings"
)

type RolesNeoEngine struct {
	Cr neocypherrunner.CypherRunner
}

func (bnc RolesNeoEngine) DecodeJSON(dec *json.Decoder) (interface{}, string, error) {
	b := Organisation{}
	err := dec.Decode(&b)
	return b, b.UUID, err
}

func (bnc RolesNeoEngine) SuggestedIndexes() map[string]string {
	return map[string]string{
		"Organisation": "uuid",
		"Concept":      "uuid",
		"Thing":        "uuid",
	}
}

func (bnc RolesNeoEngine) Read(identity string) (interface{}, bool, error) {
	panic("not implemented")
}

func (bnc RolesNeoEngine) CreateOrUpdate(obj interface{}) error {
	o := obj.(Organisation)

	p := map[string]interface{}{
		"uuid": o.UUID,
	}

	if o.Extinct == true {
		p["extinct"] = true
	}
	if o.FormerNames != nil && len(o.FormerNames) != 0 {
		p["formerNames"] = o.FormerNames
	}
	if o.HiddenLabel != "" {
		p["hiddenLabel"] = o.HiddenLabel
	}
	if o.LegalName != "" {
		p["legalName"] = o.LegalName
	}
	if o.LocalNames != nil && len(o.LocalNames) != 0 {
		p["localNames"] = o.LocalNames
	}
	if o.ProperName != "" {
		p["properName"] = o.ProperName
		p["prefLabel"] = o.ProperName
	}
	if o.ShortName != "" {
		p["shortName"] = o.ShortName
	}
	if o.TradeNames != nil && len(o.TradeNames) != 0 {
		p["tradeNames"] = o.TradeNames
	}
	for _, identifier := range o.Identifiers {
		if identifier.Authority == fsAuthority {
			p["factsetIdentifier"] = identifier.IdentifierValue
		}
		if identifier.Authority == leiAuthority {
			p["leiIdentifier"] = identifier.IdentifierValue
		}
	}
	p["uuid"] = o.UUID

	parms := map[string]interface{}{
		"uuid":     o.UUID,
		"allProps": neoism.Props(p),
		"puuid":    o.ParentOrganisation,
		"icuuid":   o.IndustryClassification,
	}

	statement := `
			MERGE (n:Thing {uuid: {uuid}})
			SET n = {allProps}
			SET n :Concept
			SET n :Organisation
		`

	nextType := o.Type
	for nextType != "Organisation" && nextType != "" {
		statement += fmt.Sprintf("SET n :%s\n", nextType)
		nextType = superTypes[nextType]
	}

	if o.ParentOrganisation != "" {
		statement += `
			MERGE (p:Thing {uuid: {puuid}})
			MERGE (n)-[:SUB_ORG_OF]->(p)
		`
	}

	if o.IndustryClassification != "" {
		statement += `
			MERGE (ic:Thing {uuid: {icuuid}})
			MERGE (n)-[:IN_INDUSTRY]->(ic)
		`
	}

	queries := []*neoism.CypherQuery{
		&neoism.CypherQuery{Statement: statement, Parameters: parms},
	}

	return bnc.Cr.CypherBatch(queries)
}

func (bnc RolesNeoEngine) Delete(identity string) (bool, error) {
	panic("not implemented")
}

func uriToUUID(uri string) string {
	// TODO: make this more robust
	return strings.Replace(uri, "http://api.ft.com/things/", "", 1)
}

const (
	fsAuthority  = "http://api.ft.com/system/FACTSET-EDM"
	leiAuthority = "http://api.ft.com/system/LEI"
)

var superTypes = map[string]string{
	"Organisation":  "",
	"Company":       "Organisation",
	"PublicCompany": "Company",
}
<commit_msg>update to match latest interface<commit_after>package orgs

import (
	"encoding/json"
	"fmt"
	"github.com/Financial-Times/neo-cypher-runner-go"
	"github.com/jmcvetta/neoism"
	"strings"
)

type RolesNeoEngine struct {
	Cr neocypherrunner.CypherRunner
}

func (bnc RolesNeoEngine) DecodeJSON(dec *json.Decoder) (interface{}, string, error) {
	b := Organisation{}
	err := dec.Decode(&b)
	return b, b.UUID, err
}

func (bnc RolesNeoEngine) SuggestedIndexes() map[string]string {
	return map[string]string{
		"Organisation": "uuid",
		"Concept":      "uuid",
		"Thing":        "uuid",
	}
}

func (bnc RolesNeoEngine) Read(identity string) (interface{}, bool, error) {
	panic("not implemented")
}

func (bnc RolesNeoEngine) Write(obj interface{}) error {
	o := obj.(Organisation)

	p := map[string]interface{}{
		"uuid": o.UUID,
	}

	if o.Extinct == true {
		p["extinct"] = true
	}
	if o.FormerNames != nil && len(o.FormerNames) != 0 {
		p["formerNames"] = o.FormerNames
	}
	if o.HiddenLabel != "" {
		p["hiddenLabel"] = o.HiddenLabel
	}
	if o.LegalName != "" {
		p["legalName"] = o.LegalName
	}
	if o.LocalNames != nil && len(o.LocalNames) != 0 {
		p["localNames"] = o.LocalNames
	}
	if o.ProperName != "" {
		p["properName"] = o.ProperName
		p["prefLabel"] = o.ProperName
	}
	if o.ShortName != "" {
		p["shortName"] = o.ShortName
	}
	if o.TradeNames != nil && len(o.TradeNames) != 0 {
		p["tradeNames"] = o.TradeNames
	}
	for _, identifier := range o.Identifiers {
		if identifier.Authority == fsAuthority {
			p["factsetIdentifier"] = identifier.IdentifierValue
		}
		if identifier.Authority == leiAuthority {
			p["leiIdentifier"] = identifier.IdentifierValue
		}
	}
	p["uuid"] = o.UUID

	parms := map[string]interface{}{
		"uuid":     o.UUID,
		"allProps": neoism.Props(p),
		"puuid":    o.ParentOrganisation,
		"icuuid":   o.IndustryClassification,
	}

	statement := `
			MERGE (n:Thing {uuid: {uuid}})
			SET n = {allProps}
			SET n :Concept
			SET n :Organisation
		`

	nextType := o.Type
	for nextType != "Organisation" && nextType != "" {
		statement += fmt.Sprintf("SET n :%s\n", nextType)
		nextType = superTypes[nextType]
	}

	if o.ParentOrganisation != "" {
		statement += `
			MERGE (p:Thing {uuid: {puuid}})
			MERGE (n)-[:SUB_ORG_OF]->(p)
		`
	}

	if o.IndustryClassification != "" {
		statement += `
			MERGE (ic:Thing {uuid: {icuuid}})
			MERGE (n)-[:IN_INDUSTRY]->(ic)
		`
	}

	queries := []*neoism.CypherQuery{
		&neoism.CypherQuery{Statement: statement, Parameters: parms},
	}

	return bnc.Cr.CypherBatch(queries)
}

func (bnc RolesNeoEngine) Delete(identity string) (bool, error) {
	panic("not implemented")
}

func uriToUUID(uri string) string {
	// TODO: make this more robust
	return strings.Replace(uri, "http://api.ft.com/things/", "", 1)
}

const (
	fsAuthority  = "http://api.ft.com/system/FACTSET-EDM"
	leiAuthority = "http://api.ft.com/system/LEI"
)

var superTypes = map[string]string{
	"Organisation":  "",
	"Company":       "Organisation",
	"PublicCompany": "Company",
}
<|endoftext|>


<commit_before>/*-
 * Copyright 2016 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package lib

import (
	"bufio"
	"bytes"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"text/template"
	"time"

	"crypto/x509/pkix"

	"github.com/fatih/color"
)

var verboseLayout = `
{{- define "PkixName" -}}
{{- range .Names}}
	{{ .Type | oidName }}: {{ .Value }}
{{- end -}}
{{end -}}

{{- if .Alias}}{{.Alias}}
{{end}}Serial: {{.SerialNumber}}
Not Before: {{.NotBefore | certStart}}
Not After : {{.NotAfter | certEnd}}
Signature : {{.SignatureAlgorithm | highlightAlgorithm}}{{if .IsSelfSigned}} (self-signed){{end}}
Subject Info:
	{{- template "PkixName" .Subject.Name}}
Issuer Info:
	{{- template "PkixName" .Issuer.Name}}
{{- if .Subject.KeyID}}
Subject Key ID   : {{.Subject.KeyID | hexify}}{{end}}{{if .Issuer.KeyID}}
Authority Key ID : {{.Issuer.KeyID | hexify}}{{end}}{{if .BasicConstraints}}
Basic Constraints: CA:{{.BasicConstraints.IsCA}}{{if .BasicConstraints.MaxPathLen}}, pathlen:{{.BasicConstraints.MaxPathLen}}{{end}}{{end}}{{if .NameConstraints}}
Name Constraints {{if .PermittedDNSDomains.Critical}}(critical){{end}}: {{range .NameConstraints.PermittedDNSDomains}}
	{{.}}{{end}}{{end}}{{if .KeyUsage}}
Key Usage:{{range .KeyUsage | keyUsage}}
	{{.}}{{end}}{{end}}{{if .ExtKeyUsage}}
Extended Key Usage:{{range .ExtKeyUsage}}
	{{. | extKeyUsage}}{{end}}{{end}}{{if .AltDNSNames}}
Alternate DNS Names:{{range .AltDNSNames}}
	{{.}}{{end}}{{end}}{{if .AltIPAddresses}}
Alternate IP Addresses:{{range .AltIPAddresses}}
	{{.}}{{end}}{{end}}{{if .EmailAddresses}}
Email Addresses:{{range .EmailAddresses}}
	{{.}}{{end}}{{end}}{{if .Warnings}}
Warnings:{{range .Warnings}}
	{{. | redify}}{{end}}{{end}}`

var layout = `
{{- if .Alias}}{{.Alias}}
{{end -}}
Valid: {{.NotBefore | certStart}} to {{.NotAfter | certEnd}}
Subject: {{.Subject.Name | printName }}
Issuer: {{.Issuer.Name | printName }}
{{- if .NameConstraints}}
Name Constraints {{if .PermittedDNSDomains.Critical}}(critical){{end}}: {{range .NameConstraints.PermittedDNSDomains}}
	{{.}}{{end}}{{end}}
{{- if .AltDNSNames}}
Alternate DNS Names:{{range .AltDNSNames}}
	{{.}}{{end}}{{end}}
{{- if .AltIPAddresses}}
Alternate IP Addresses:{{range .AltIPAddresses}}
	{{.}}{{end}}{{end}}
{{- if .EmailAddresses}}
Email Addresses:{{range .EmailAddresses}}
	{{.}}{{end}}{{end}}
{{- if .Warnings}}
Warnings:{{range .Warnings}}
	{{. | redify}}{{end}}{{end}}`

type certWithName struct {
	name string
	file string
	cert *x509.Certificate
}

func (c certWithName) MarshalJSON() ([]byte, error) {
	out := createSimpleCertificate(c.name, c.cert)
	return json.Marshal(out)
}

func createSimpleCertificateFromX509(block *pem.Block) (simpleCertificate, error) {
	raw, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return simpleCertificate{}, fmt.Errorf("error reading cert: %s", err)
	}

	cert := certWithName{cert: raw}
	if val, ok := block.Headers[nameHeader]; ok {
		cert.name = val
	}
	if val, ok := block.Headers[fileHeader]; ok {
		cert.file = val
	}

	return createSimpleCertificate(cert.name, cert.cert), nil
}

// EncodeX509ToJSON encodes an X.509 certificate into a JSON string.
func EncodeX509ToJSON(cert *x509.Certificate) []byte {
	out := createSimpleCertificate("", cert)
	raw, err := json.Marshal(out)
	if err != nil {
		panic(err)
	}
	return raw
}

// EncodeX509ToObject encodes an X.509 certificate into a JSON-serializable object.
func EncodeX509ToObject(cert *x509.Certificate) interface{} {
	return createSimpleCertificate("", cert)
}

// EncodeX509ToText encodes an X.509 certificate into human-readable text.
func EncodeX509ToText(cert *x509.Certificate, verbose bool) []byte {
	return displayCert(createSimpleCertificate("", cert), verbose)
}

// displayCert takes in a parsed certificate object
// (for jceks certs, blank otherwise), and prints out relevant
// information. Start and end dates are colored based on whether or not
// the certificate is expired, not expired, or close to expiring.
func displayCert(cert simpleCertificate, verbose bool) []byte {
	funcMap := template.FuncMap{
		"certStart":          certStart,
		"certEnd":            certEnd,
		"redify":             redify,
		"highlightAlgorithm": highlightAlgorithm,
		"hexify":             hexify,
		"keyUsage":           keyUsage,
		"extKeyUsage":        extKeyUsage,
		"oidName":            oidName,
		"oidShort":           oidShort,
		"printName":          printName,
	}
	t := template.New("Cert template").Funcs(funcMap)
	var err error
	if verbose {
		t, err = t.Parse(verboseLayout)
	} else {
		t, err = t.Parse(layout)
	}
	if err != nil {
		// Should never happen
		panic(err)
	}
	var buffer bytes.Buffer
	w := bufio.NewWriter(&buffer)
	err = t.Execute(w, cert)
	if err != nil {
		// Should never happen
		panic(err)
	}
	w.Flush()
	return buffer.Bytes()
}

var (
	green  = color.New(color.Bold, color.FgGreen)
	yellow = color.New(color.Bold, color.FgYellow)
	red    = color.New(color.Bold, color.FgRed)
)

var algorithmColors = map[x509.SignatureAlgorithm]*color.Color{
	x509.MD2WithRSA:      red,
	x509.MD5WithRSA:      red,
	x509.SHA1WithRSA:     red,
	x509.SHA256WithRSA:   green,
	x509.SHA384WithRSA:   green,
	x509.SHA512WithRSA:   green,
	x509.DSAWithSHA1:     red,
	x509.DSAWithSHA256:   red,
	x509.ECDSAWithSHA1:   red,
	x509.ECDSAWithSHA256: green,
	x509.ECDSAWithSHA384: green,
	x509.ECDSAWithSHA512: green,
}

// highlightAlgorithm changes the color of the signing algorithm
// based on a set color map, e.g. to make SHA-1 show up red.
func highlightAlgorithm(sigAlg simpleSigAlg) string {
	sig := x509.SignatureAlgorithm(sigAlg)
	color, ok := algorithmColors[sig]
	if !ok {
		return algString(sig)
	}
	return color.SprintFunc()(algString(sig))
}

// timeString formats a time in UTC with minute precision, in the given color.
func timeString(t time.Time, c *color.Color) string {
	return c.SprintfFunc()(t.Format("2006-01-02 15:04Z"))
}

// certStart takes a given start time for the validity of
// a certificate and returns that time colored properly
// based on how close it is to expiry. If it's more than
// a day after the certificate became valid the string will
// be green. If it has been less than a day the string will
// be yellow. If the certificate is not yet valid, the string
// will be red.
func certStart(start time.Time) string {
	now := time.Now()
	day, _ := time.ParseDuration("24h")
	threshold := start.Add(day)
	if now.After(threshold) {
		return timeString(start, green)
	} else if now.After(start) {
		return timeString(start, yellow)
	} else {
		return timeString(start, red)
	}
}

// certEnd takes a given end time for the validity of
// a certificate and returns that time colored properly
// based on how close it is to expiry. If the certificate
// is more than a month away from expiring it returns a
// green string. If the certificate is less than a month
// from expiry it returns a yellow string. If the certificate
// is expired it returns a red string.
func certEnd(end time.Time) string {
	now := time.Now()
	month, _ := time.ParseDuration("720h")
	threshold := now.Add(month)
	if threshold.Before(end) {
		return timeString(end, green)
	} else if now.Before(end) {
		return timeString(end, yellow)
	} else {
		return timeString(end, red)
	}
}

func redify(text string) string {
	return red.SprintfFunc()("%s", text)
}

func printName(name pkix.Name) (out string) {
	printed := false
	for _, name := range name.Names {
		short := oidShort(name.Type)
		if short != "" {
			if printed {
				out += ", "
			}
			out += fmt.Sprintf("%s=%v", short, name.Value)
			printed = true
		}
	}

	return
}
<commit_msg>Use three-letter timezone instead of fixed 'Z'<commit_after>/*-
 * Copyright 2016 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package lib

import (
	"bufio"
	"bytes"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"text/template"
	"time"

	"crypto/x509/pkix"

	"github.com/fatih/color"
)

var verboseLayout = `
{{- define "PkixName" -}}
{{- range .Names}}
	{{ .Type | oidName }}: {{ .Value }}
{{- end -}}
{{end -}}

{{- if .Alias}}{{.Alias}}
{{end}}Serial: {{.SerialNumber}}
Not Before: {{.NotBefore | certStart}}
Not After : {{.NotAfter | certEnd}}
Signature : {{.SignatureAlgorithm | highlightAlgorithm}}{{if .IsSelfSigned}} (self-signed){{end}}
Subject Info:
	{{- template "PkixName" .Subject.Name}}
Issuer Info:
	{{- template "PkixName" .Issuer.Name}}
{{- if .Subject.KeyID}}
Subject Key ID   : {{.Subject.KeyID | hexify}}{{end}}{{if .Issuer.KeyID}}
Authority Key ID : {{.Issuer.KeyID | hexify}}{{end}}{{if .BasicConstraints}}
Basic Constraints: CA:{{.BasicConstraints.IsCA}}{{if .BasicConstraints.MaxPathLen}}, pathlen:{{.BasicConstraints.MaxPathLen}}{{end}}{{end}}{{if .NameConstraints}}
Name Constraints {{if .PermittedDNSDomains.Critical}}(critical){{end}}: {{range .NameConstraints.PermittedDNSDomains}}
	{{.}}{{end}}{{end}}{{if .KeyUsage}}
Key Usage:{{range .KeyUsage | keyUsage}}
	{{.}}{{end}}{{end}}{{if .ExtKeyUsage}}
Extended Key Usage:{{range .ExtKeyUsage}}
	{{. | extKeyUsage}}{{end}}{{end}}{{if .AltDNSNames}}
Alternate DNS Names:{{range .AltDNSNames}}
	{{.}}{{end}}{{end}}{{if .AltIPAddresses}}
Alternate IP Addresses:{{range .AltIPAddresses}}
	{{.}}{{end}}{{end}}{{if .EmailAddresses}}
Email Addresses:{{range .EmailAddresses}}
	{{.}}{{end}}{{end}}{{if .Warnings}}
Warnings:{{range .Warnings}}
	{{. | redify}}{{end}}{{end}}`

var layout = `
{{- if .Alias}}{{.Alias}}
{{end -}}
Valid: {{.NotBefore | certStart}} to {{.NotAfter | certEnd}}
Subject: {{.Subject.Name | printName }}
Issuer: {{.Issuer.Name | printName }}
{{- if .NameConstraints}}
Name Constraints {{if .PermittedDNSDomains.Critical}}(critical){{end}}: {{range .NameConstraints.PermittedDNSDomains}}
	{{.}}{{end}}{{end}}
{{- if .AltDNSNames}}
Alternate DNS Names:{{range .AltDNSNames}}
	{{.}}{{end}}{{end}}
{{- if .AltIPAddresses}}
Alternate IP Addresses:{{range .AltIPAddresses}}
	{{.}}{{end}}{{end}}
{{- if .EmailAddresses}}
Email Addresses:{{range .EmailAddresses}}
	{{.}}{{end}}{{end}}
{{- if .Warnings}}
Warnings:{{range .Warnings}}
	{{. | redify}}{{end}}{{end}}`

type certWithName struct {
	name string
	file string
	cert *x509.Certificate
}

func (c certWithName) MarshalJSON() ([]byte, error) {
	out := createSimpleCertificate(c.name, c.cert)
	return json.Marshal(out)
}

func createSimpleCertificateFromX509(block *pem.Block) (simpleCertificate, error) {
	raw, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return simpleCertificate{}, fmt.Errorf("error reading cert: %s", err)
	}

	cert := certWithName{cert: raw}
	if val, ok := block.Headers[nameHeader]; ok {
		cert.name = val
	}
	if val, ok := block.Headers[fileHeader]; ok {
		cert.file = val
	}

	return createSimpleCertificate(cert.name, cert.cert), nil
}

// EncodeX509ToJSON encodes an X.509 certificate into a JSON string.
func EncodeX509ToJSON(cert *x509.Certificate) []byte {
	out := createSimpleCertificate("", cert)
	raw, err := json.Marshal(out)
	if err != nil {
		panic(err)
	}
	return raw
}

// EncodeX509ToObject encodes an X.509 certificate into a JSON-serializable object.
func EncodeX509ToObject(cert *x509.Certificate) interface{} {
	return createSimpleCertificate("", cert)
}

// EncodeX509ToText encodes an X.509 certificate into human-readable text.
func EncodeX509ToText(cert *x509.Certificate, verbose bool) []byte {
	return displayCert(createSimpleCertificate("", cert), verbose)
}

// displayCert takes in a parsed certificate object
// (for jceks certs, blank otherwise), and prints out relevant
// information. Start and end dates are colored based on whether or not
// the certificate is expired, not expired, or close to expiring.
func displayCert(cert simpleCertificate, verbose bool) []byte {
	funcMap := template.FuncMap{
		"certStart":          certStart,
		"certEnd":            certEnd,
		"redify":             redify,
		"highlightAlgorithm": highlightAlgorithm,
		"hexify":             hexify,
		"keyUsage":           keyUsage,
		"extKeyUsage":        extKeyUsage,
		"oidName":            oidName,
		"oidShort":           oidShort,
		"printName":          printName,
	}
	t := template.New("Cert template").Funcs(funcMap)
	var err error
	if verbose {
		t, err = t.Parse(verboseLayout)
	} else {
		t, err = t.Parse(layout)
	}
	if err != nil {
		// Should never happen
		panic(err)
	}
	var buffer bytes.Buffer
	w := bufio.NewWriter(&buffer)
	err = t.Execute(w, cert)
	if err != nil {
		// Should never happen
		panic(err)
	}
	w.Flush()
	return buffer.Bytes()
}

var (
	green  = color.New(color.Bold, color.FgGreen)
	yellow = color.New(color.Bold, color.FgYellow)
	red    = color.New(color.Bold, color.FgRed)
)

var algorithmColors = map[x509.SignatureAlgorithm]*color.Color{
	x509.MD2WithRSA:      red,
	x509.MD5WithRSA:      red,
	x509.SHA1WithRSA:     red,
	x509.SHA256WithRSA:   green,
	x509.SHA384WithRSA:   green,
	x509.SHA512WithRSA:   green,
	x509.DSAWithSHA1:     red,
	x509.DSAWithSHA256:   red,
	x509.ECDSAWithSHA1:   red,
	x509.ECDSAWithSHA256: green,
	x509.ECDSAWithSHA384: green,
	x509.ECDSAWithSHA512: green,
}

// highlightAlgorithm changes the color of the signing algorithm
// based on a set color map, e.g. to make SHA-1 show up red.
func highlightAlgorithm(sigAlg simpleSigAlg) string {
	sig := x509.SignatureAlgorithm(sigAlg)
	color, ok := algorithmColors[sig]
	if !ok {
		return algString(sig)
	}
	return color.SprintFunc()(algString(sig))
}

// timeString formats a time in UTC with minute precision, in the given color.
func timeString(t time.Time, c *color.Color) string {
	return c.SprintfFunc()(t.Format("2006-01-02 15:04 MST"))
}

// certStart takes a given start time for the validity of
// a certificate and returns that time colored properly
// based on how close it is to expiry. If it's more than
// a day after the certificate became valid the string will
// be green. If it has been less than a day the string will
// be yellow. If the certificate is not yet valid, the string
// will be red.
func certStart(start time.Time) string {
	now := time.Now()
	day, _ := time.ParseDuration("24h")
	threshold := start.Add(day)
	if now.After(threshold) {
		return timeString(start, green)
	} else if now.After(start) {
		return timeString(start, yellow)
	} else {
		return timeString(start, red)
	}
}

// certEnd takes a given end time for the validity of
// a certificate and returns that time colored properly
// based on how close it is to expiry. If the certificate
// is more than a month away from expiring it returns a
// green string. If the certificate is less than a month
// from expiry it returns a yellow string. If the certificate
// is expired it returns a red string.
func certEnd(end time.Time) string {
	now := time.Now()
	month, _ := time.ParseDuration("720h")
	threshold := now.Add(month)
	if threshold.Before(end) {
		return timeString(end, green)
	} else if now.Before(end) {
		return timeString(end, yellow)
	} else {
		return timeString(end, red)
	}
}

func redify(text string) string {
	return red.SprintfFunc()("%s", text)
}

func printName(name pkix.Name) (out string) {
	printed := false
	for _, name := range name.Names {
		short := oidShort(name.Type)
		if short != "" {
			if printed {
				out += ", "
			}
			out += fmt.Sprintf("%s=%v", short, name.Value)
			printed = true
		}
	}

	return
}
<|endoftext|>


<commit_before>package stage0

//
// Rocket is a reference implementation of the app container specification.
//
// Execution on Rocket is divided into a number of stages, and the `rkt`
// binary implements the first stage (stage 0)
//

import (
	"archive/tar"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"syscall"

	"github.com/coreos/rocket/Godeps/_workspace/src/code.google.com/p/go-uuid/uuid"
	"github.com/coreos/rocket/app-container/aci"
	"github.com/coreos/rocket/app-container/schema"
	"github.com/coreos/rocket/app-container/schema/types"
	"github.com/coreos/rocket/cas"
	rktpath "github.com/coreos/rocket/path"
	ptar "github.com/coreos/rocket/pkg/tar"
	"github.com/coreos/rocket/version"

	"github.com/coreos/rocket/stage0/stage1_init"
	"github.com/coreos/rocket/stage0/stage1_rootfs"
)

const (
	initPath = "stage1/init"
)

type Config struct {
	Store         *cas.Store // store containing all of the configured application images
	ContainersDir string     // root directory for rocket containers
	Stage1Init    string     // binary to be execed as stage1
	Stage1Rootfs  string     // compressed bundle containing a rootfs for stage1
	Debug         bool
	Images        []types.Hash      // application images
	Volumes       map[string]string // map of volumes that rocket can provide to applications
}

func init() {
	log.SetOutput(ioutil.Discard)
}

// Setup sets up a filesystem for a container based on the given config.
// The directory containing the filesystem is returned, and any error encountered.
func Setup(cfg Config) (string, error) {
	if cfg.Debug {
		log.SetOutput(os.Stderr)
	}

	cuuid, err := types.NewUUID(uuid.New())
	if err != nil {
		return "", fmt.Errorf("error creating UID: %v", err)
	}

	// TODO(jonboulle): collision detection/mitigation
	// Create a directory for this container
	dir := filepath.Join(cfg.ContainersDir, cuuid.String())

	if err := os.MkdirAll(dir, 0700); err != nil {
		return "", fmt.Errorf("error creating directory: %v", err)
	}

	log.Printf("Unpacking stage1 rootfs")
	if cfg.Stage1Rootfs != "" {
		if err = unpackRootfs(cfg.Stage1Rootfs, rktpath.Stage1RootfsPath(dir)); err != nil {
			return "", fmt.Errorf("error unpacking rootfs: %v", err)
		}
	} else {
		if err = unpackBuiltinRootfs(rktpath.Stage1RootfsPath(dir)); err != nil {
			return "", fmt.Errorf("error unpacking rootfs: %v", err)
		}
	}

	log.Printf("Writing stage1 init")
	var in io.Reader
	if cfg.Stage1Init != "" {
		in, err = os.Open(cfg.Stage1Init)
		if err != nil {
			return "", fmt.Errorf("error loading stage1 init binary: %v", err)
		}
	} else {
		init_bin, err := stage1_init.Asset("s1init")
		if err != nil {
			return "", fmt.Errorf("error accessing stage1 init bindata: %v", err)
		}
		in = bytes.NewBuffer(init_bin)
	}
	fn := filepath.Join(dir, initPath)
	out, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY, 0555)
	if err != nil {
		return "", fmt.Errorf("error opening stage1 init for writing: %v", err)
	}
	if _, err := io.Copy(out, in); err != nil {
		return "", fmt.Errorf("error writing stage1 init: %v", err)
	}
	if err := out.Close(); err != nil {
		return "", fmt.Errorf("error closing stage1 init: %v", err)
	}

	log.Printf("Wrote filesystem to %s\n", dir)

	cm := schema.ContainerRuntimeManifest{
		ACKind: "ContainerRuntimeManifest",
		UUID:   *cuuid,
		Apps:   make(schema.AppList, 0),
	}

	v, err := types.NewSemVer(version.Version)
	if err != nil {
		return "", fmt.Errorf("error creating version: %v", err)
	}
	cm.ACVersion = *v

	for _, img := range cfg.Images {
		am, err := setupImage(cfg, img, dir)
		if err != nil {
			return "", fmt.Errorf("error setting up image %s: %v", img, err)
		}
		if cm.Apps.Get(am.Name) != nil {
			return "", fmt.Errorf("error: multiple apps with name %s", am.Name)
		}
		a := schema.App{
			Name:        am.Name,
			ImageID:     img,
			Isolators:   am.Isolators,
			Annotations: am.Annotations,
		}
		cm.Apps = append(cm.Apps, a)
	}

	var sVols []types.Volume
	for key, path := range cfg.Volumes {
		v := types.Volume{
			Kind:     "host",
			Source:   path,
			ReadOnly: true,
			Fulfills: []types.ACName{
				types.ACName(key),
			},
		}
		sVols = append(sVols, v)
	}
	// TODO(jonboulle): check that app mountpoint expectations are
	// satisfied here, rather than waiting for stage1
	cm.Volumes = sVols

	cdoc, err := json.Marshal(cm)
	if err != nil {
		return "", fmt.Errorf("error marshalling container manifest: %v", err)
	}

	log.Printf("Writing container manifest")
	fn = rktpath.ContainerManifestPath(dir)
	if err := ioutil.WriteFile(fn, cdoc, 0700); err != nil {
		return "", fmt.Errorf("error writing container manifest: %v", err)
	}
	return dir, nil
}

// Run actually runs the container by exec()ing the stage1 init inside
// the container filesystem.
func Run(dir string, debug bool) {
	log.Printf("Pivoting to filesystem %s", dir)
	if err := os.Chdir(dir); err != nil {
		log.Fatalf("failed changing to dir: %v", err)
	}

	log.Printf("Execing %s", initPath)
	args := []string{initPath}
	if debug {
		args = append(args, "debug")
	}
	if err := syscall.Exec(initPath, args, os.Environ()); err != nil {
		log.Fatalf("error execing init: %v", err)
	}
}

func untarRootfs(r io.Reader, dir string) error {
	tr := tar.NewReader(r)
	if err := os.MkdirAll(dir, 0776); err != nil {
		return fmt.Errorf("error creating stage1 rootfs directory: %v", err)
	}

	if err := ptar.ExtractTar(tr, dir); err != nil {
		return fmt.Errorf("error extracting rootfs: %v", err)
	}
	return nil
}

// unpackRootfs unpacks a stage1 rootfs (compressed file, pointed to by rfs)
// into dir, returning any error encountered
func unpackRootfs(rfs string, dir string) error {
	fh, err := os.Open(rfs)
	if err != nil {
		return fmt.Errorf("error opening stage1 rootfs: %v", err)
	}
	typ, err := aci.DetectFileType(fh)
	if err != nil {
		return fmt.Errorf("error detecting image type: %v", err)
	}
	if _, err := fh.Seek(0, 0); err != nil {
		return fmt.Errorf("error seeking image: %v", err)
	}
	var r io.Reader
	switch typ {
	case aci.TypeGzip:
		r, err = gzip.NewReader(fh)
		if err != nil {
			return fmt.Errorf("error reading gzip: %v", err)
		}
	case aci.TypeBzip2:
		r = bzip2.NewReader(fh)
	case aci.TypeXz:
		r = aci.XzReader(fh)
	case aci.TypeUnknown:
		return fmt.Errorf("error: unknown image filetype")
	default:
		// should never happen
		panic("no type returned from DetectFileType?")
	}

	if err := untarRootfs(r, dir); err != nil {
		return fmt.Errorf("error untarring rootfs")
	}

	return nil
}

// unpackBuiltinRootfs unpacks the included stage1 rootfs into dir
func unpackBuiltinRootfs(dir string) error {
	b, err := stage1_rootfs.Asset("s1rootfs.tar")
	if err != nil {
		return fmt.Errorf("error accessing rootfs asset: %v", err)
	}
	buf := bytes.NewBuffer(b)

	if err = untarRootfs(buf, dir); err != nil {
		return fmt.Errorf("error untarring rootfs")
	}

	return nil
}

// setupImage attempts to load the image by the given hash from the store,
// verifies that the image matches the given hash and extracts the image
// into a directory in the given dir.
// It returns the AppManifest that the image contains
func setupImage(cfg Config, img types.Hash, dir string) (*schema.AppManifest, error) {
	log.Println("Loading image", img.String())

	rs, err := cfg.Store.ReadStream(img.String())
	if err != nil {
		return nil, err
	}

	ad := rktpath.AppImagePath(dir, img)
	err = os.MkdirAll(ad, 0776)
	if err != nil {
		return nil, fmt.Errorf("error creating image directory: %v", err)
	}

	hash := sha256.New()
	r := io.TeeReader(rs, hash)

	if err := ptar.ExtractTar(tar.NewReader(r), ad); err != nil {
		return nil, fmt.Errorf("error extracting ACI: %v", err)
	}

	// Tar does not necessarily read the complete file, so ensure we read the entirety into the hash
	if _, err := io.Copy(ioutil.Discard, r); err != nil {
		return nil, fmt.Errorf("error reading ACI: %v", err)
	}

	if id := fmt.Sprintf("%x", hash.Sum(nil)); id != img.Val {
		if err := os.RemoveAll(ad); err != nil {
			fmt.Fprintf(os.Stderr, "error cleaning up directory: %v\n", err)
		}
		return nil, fmt.Errorf("image hash does not match expected (%v != %v)", id, img.Val)
	}

	err = os.MkdirAll(filepath.Join(ad, "rootfs/tmp"), 0777)
	if err != nil {
		return nil, fmt.Errorf("error creating tmp directory: %v", err)
	}

	mpath := rktpath.AppManifestPath(dir, img)
	f, err := os.Open(mpath)
	if err != nil {
		return nil, fmt.Errorf("error opening app manifest: %v", err)
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, fmt.Errorf("error reading app manifest: %v", err)
	}
	var am schema.AppManifest
	if err := json.Unmarshal(b, &am); err != nil {
		return nil, fmt.Errorf("error unmarshaling app manifest: %v", err)
	}
	return &am, nil
}
<commit_msg>stage0: clear up rootfs unpacking error propagation<commit_after>package stage0

//
// Rocket is a reference implementation of the app container specification.
//
// Execution on Rocket is divided into a number of stages, and the `rkt`
// binary implements the first stage (stage 0)
//

import (
	"archive/tar"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"syscall"

	"github.com/coreos/rocket/Godeps/_workspace/src/code.google.com/p/go-uuid/uuid"
	"github.com/coreos/rocket/app-container/aci"
	"github.com/coreos/rocket/app-container/schema"
	"github.com/coreos/rocket/app-container/schema/types"
	"github.com/coreos/rocket/cas"
	rktpath "github.com/coreos/rocket/path"
	ptar "github.com/coreos/rocket/pkg/tar"
	"github.com/coreos/rocket/version"

	"github.com/coreos/rocket/stage0/stage1_init"
	"github.com/coreos/rocket/stage0/stage1_rootfs"
)

const (
	initPath = "stage1/init"
)

type Config struct {
	Store         *cas.Store // store containing all of the configured application images
	ContainersDir string     // root directory for rocket containers
	Stage1Init    string     // binary to be execed as stage1
	Stage1Rootfs  string     // compressed bundle containing a rootfs for stage1
	Debug         bool
	Images        []types.Hash      // application images
	Volumes       map[string]string // map of volumes that rocket can provide to applications
}

func init() {
	log.SetOutput(ioutil.Discard)
}

// Setup sets up a filesystem for a container based on the given config.
// The directory containing the filesystem is returned, and any error encountered.
func Setup(cfg Config) (string, error) {
	if cfg.Debug {
		log.SetOutput(os.Stderr)
	}

	cuuid, err := types.NewUUID(uuid.New())
	if err != nil {
		return "", fmt.Errorf("error creating UID: %v", err)
	}

	// TODO(jonboulle): collision detection/mitigation
	// Create a directory for this container
	dir := filepath.Join(cfg.ContainersDir, cuuid.String())

	if err := os.MkdirAll(dir, 0700); err != nil {
		return "", fmt.Errorf("error creating directory: %v", err)
	}

	log.Printf("Unpacking stage1 rootfs")
	if cfg.Stage1Rootfs != "" {
		err = unpackRootfs(cfg.Stage1Rootfs, rktpath.Stage1RootfsPath(dir))
	} else {
		err = unpackBuiltinRootfs(rktpath.Stage1RootfsPath(dir))
	}
	if err != nil {
		return "", fmt.Errorf("error unpacking rootfs: %v", err)
	}

	log.Printf("Writing stage1 init")
	var in io.Reader
	if cfg.Stage1Init != "" {
		in, err = os.Open(cfg.Stage1Init)
		if err != nil {
			return "", fmt.Errorf("error loading stage1 init binary: %v", err)
		}
	} else {
		init_bin, err := stage1_init.Asset("s1init")
		if err != nil {
			return "", fmt.Errorf("error accessing stage1 init bindata: %v", err)
		}
		in = bytes.NewBuffer(init_bin)
	}
	fn := filepath.Join(dir, initPath)
	out, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY, 0555)
	if err != nil {
		return "", fmt.Errorf("error opening stage1 init for writing: %v", err)
	}
	if _, err := io.Copy(out, in); err != nil {
		return "", fmt.Errorf("error writing stage1 init: %v", err)
	}
	if err := out.Close(); err != nil {
		return "", fmt.Errorf("error closing stage1 init: %v", err)
	}

	log.Printf("Wrote filesystem to %s\n", dir)

	cm := schema.ContainerRuntimeManifest{
		ACKind: "ContainerRuntimeManifest",
		UUID:   *cuuid,
		Apps:   make(schema.AppList, 0),
	}

	v, err := types.NewSemVer(version.Version)
	if err != nil {
		return "", fmt.Errorf("error creating version: %v", err)
	}
	cm.ACVersion = *v

	for _, img := range cfg.Images {
		am, err := setupImage(cfg, img, dir)
		if err != nil {
			return "", fmt.Errorf("error setting up image %s: %v", img, err)
		}
		if cm.Apps.Get(am.Name) != nil {
			return "", fmt.Errorf("error: multiple apps with name %s", am.Name)
		}
		a := schema.App{
			Name:        am.Name,
			ImageID:     img,
			Isolators:   am.Isolators,
			Annotations: am.Annotations,
		}
		cm.Apps = append(cm.Apps, a)
	}

	var sVols []types.Volume
	for key, path := range cfg.Volumes {
		v := types.Volume{
			Kind:     "host",
			Source:   path,
			ReadOnly: true,
			Fulfills: []types.ACName{
				types.ACName(key),
			},
		}
		sVols = append(sVols, v)
	}
	// TODO(jonboulle): check that app mountpoint expectations are
	// satisfied here, rather than waiting for stage1
	cm.Volumes = sVols

	cdoc, err := json.Marshal(cm)
	if err != nil {
		return "", fmt.Errorf("error marshalling container manifest: %v", err)
	}

	log.Printf("Writing container manifest")
	fn = rktpath.ContainerManifestPath(dir)
	if err := ioutil.WriteFile(fn, cdoc, 0700); err != nil {
		return "", fmt.Errorf("error writing container manifest: %v", err)
	}
	return dir, nil
}

// Run actually runs the container by exec()ing the stage1 init inside
// the container filesystem.
func Run(dir string, debug bool) {
	log.Printf("Pivoting to filesystem %s", dir)
	if err := os.Chdir(dir); err != nil {
		log.Fatalf("failed changing to dir: %v", err)
	}

	log.Printf("Execing %s", initPath)
	args := []string{initPath}
	if debug {
		args = append(args, "debug")
	}
	if err := syscall.Exec(initPath, args, os.Environ()); err != nil {
		log.Fatalf("error execing init: %v", err)
	}
}

func untarRootfs(r io.Reader, dir string) error {
	tr := tar.NewReader(r)
	if err := os.MkdirAll(dir, 0776); err != nil {
		return fmt.Errorf("error creating stage1 rootfs directory: %v", err)
	}

	if err := ptar.ExtractTar(tr, dir); err != nil {
		return fmt.Errorf("error extracting rootfs: %v", err)
	}
	return nil
}

// unpackRootfs unpacks a stage1 rootfs (compressed file, pointed to by rfs)
// into dir, returning any error encountered
func unpackRootfs(rfs string, dir string) error {
	fh, err := os.Open(rfs)
	if err != nil {
		return fmt.Errorf("error opening stage1 rootfs: %v", err)
	}
	typ, err := aci.DetectFileType(fh)
	if err != nil {
		return fmt.Errorf("error detecting image type: %v", err)
	}
	if _, err := fh.Seek(0, 0); err != nil {
		return fmt.Errorf("error seeking image: %v", err)
	}
	var r io.Reader
	switch typ {
	case aci.TypeGzip:
		r, err = gzip.NewReader(fh)
		if err != nil {
			return fmt.Errorf("error reading gzip: %v", err)
		}
	case aci.TypeBzip2:
		r = bzip2.NewReader(fh)
	case aci.TypeXz:
		r = aci.XzReader(fh)
	case aci.TypeUnknown:
		return fmt.Errorf("error: unknown image filetype")
	default:
		// should never happen
		panic("no type returned from DetectFileType?")
	}
	return untarRootfs(r, dir)
}

// unpackBuiltinRootfs unpacks the included stage1 rootfs into dir
func unpackBuiltinRootfs(dir string) error {
	b, err := stage1_rootfs.Asset("s1rootfs.tar")
	if err != nil {
		return fmt.Errorf("error accessing rootfs asset: %v", err)
	}
	buf := bytes.NewBuffer(b)
	return untarRootfs(buf, dir)
}

// setupImage attempts to load the image by the given hash from the store,
// verifies that the image matches the given hash and extracts the image
// into a directory in the given dir.
// It returns the AppManifest that the image contains
func setupImage(cfg Config, img types.Hash, dir string) (*schema.AppManifest, error) {
	log.Println("Loading image", img.String())

	rs, err := cfg.Store.ReadStream(img.String())
	if err != nil {
		return nil, err
	}

	ad := rktpath.AppImagePath(dir, img)
	err = os.MkdirAll(ad, 0776)
	if err != nil {
		return nil, fmt.Errorf("error creating image directory: %v", err)
	}

	hash := sha256.New()
	r := io.TeeReader(rs, hash)

	if err := ptar.ExtractTar(tar.NewReader(r), ad); err != nil {
		return nil, fmt.Errorf("error extracting ACI: %v", err)
	}

	// Tar does not necessarily read the complete file, so ensure we read the entirety into the hash
	if _, err := io.Copy(ioutil.Discard, r); err != nil {
		return nil, fmt.Errorf("error reading ACI: %v", err)
	}

	if id := fmt.Sprintf("%x", hash.Sum(nil)); id != img.Val {
		if err := os.RemoveAll(ad); err != nil {
			fmt.Fprintf(os.Stderr, "error cleaning up directory: %v\n", err)
		}
		return nil, fmt.Errorf("image hash does not match expected (%v != %v)", id, img.Val)
	}

	err = os.MkdirAll(filepath.Join(ad, "rootfs/tmp"), 0777)
	if err != nil {
		return nil, fmt.Errorf("error creating tmp directory: %v", err)
	}

	mpath := rktpath.AppManifestPath(dir, img)
	f, err := os.Open(mpath)
	if err != nil {
		return nil, fmt.Errorf("error opening app manifest: %v", err)
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, fmt.Errorf("error reading app manifest: %v", err)
	}
	var am schema.AppManifest
	if err := json.Unmarshal(b, &am); err != nil {
		return nil, fmt.Errorf("error unmarshaling app manifest: %v", err)
	}
	return &am, nil
}
<|endoftext|>


<commit_before>package state

import (
	"errors"
	"fmt"
	"launchpad.net/gozk/zookeeper"
	"strings"
)

// Info encapsulates information about a juju state server
// or servers and can be used to make a connection to
// that state.
type Info struct {
	// Addrs gives the addresses of the zookeeper
	// servers for the state. Each address should be in the form
	// expected by net.Dial (i.e. address:port).
	Addrs []string
}

// Open connects to the server described by the given
// info and returns a new State representing the environment
// connected to.
func Open(info *Info) (*State, error) {
	if len(info.Addrs) == 0 {
		return nil, fmt.Errorf("no zookeeper addresses")
	}
	zk, session, err := zookeeper.Dial(strings.Join(info.Addrs, ","), 5e9)
	if err != nil {
		return nil, err
	}
	if !(<-session).Ok() {
		return nil, errors.New("Could not connect to zookeeper")
	}
	// throw away subsequent session events.
	go func() {
		for _ = range session {
		}
	}()
	return &State{zk}, nil
}
<commit_msg>minor changes for review<commit_after>package state

import (
	"errors"
	"fmt"
	"launchpad.net/gozk/zookeeper"
	"strings"
)

// Info encapsulates information about cluster of
// servers holding juju state and can be used to make a
// connection to that cluster.
type Info struct {
	// Addrs gives the addresses of the Zookeeper
	// servers for the state. Each address should be in the form
	// address:port.
	Addrs []string
}

const zkTimeout = 15e9

// Open connects to the server described by the given
// info and returns a new State representing the environment
// connected to.
func Open(info *Info) (*State, error) {
	if len(info.Addrs) == 0 {
		return nil, fmt.Errorf("no zookeeper addresses")
	}
	zk, session, err := zookeeper.Dial(strings.Join(info.Addrs, ","), zkTimeout)
	if err != nil {
		return nil, err
	}
	if !(<-session).Ok() {
		return nil, errors.New("Could not connect to zookeeper")
	}

	// TODO decide what to do with session events - currently
	// we will panic if the session event channel fills up.
	return &State{zk}, nil
}
<|endoftext|>


<commit_before>package sensor

import (
	"math/rand"
	"sort"
	"testing"
)

func normDistSlice(size int, stddev, mean float64) []float64 {
	res := make([]float64, 0, size)
	for i := 0; i < size; i++ {
		res = append(res, rand.NormFloat64()*stddev+mean)
	}
	return res
}

func percentile(pct float64, samples []float64) float64 {
	s := sort.Float64Slice(samples)
	s.Sort()
	return s[int(float64(len(s))*pct)]
}

func Test99Percentile(t *testing.T) {
	z99 := 2.33
	e := 0.1
	p := percentile(0.99, normDistSlice(100000, 1.0, 0.0))

	if p+e < z99 || p-e > z99 {
		t.Fatalf("99th percentile (%f) of normal distribution doesn't match the z-score of %f within an error %f", p, z99, e)
	}
}
<commit_msg>Correct implementation of quantile estimator.<commit_after>package sensor

import (
	"container/list"
	"math/rand"
	"sort"
	"testing"
)

func normDistSlice(size int, stddev, mean float64) []float64 {
	res := make([]float64, 0, size)
	for i := 0; i < size; i++ {
		res = append(res, rand.NormFloat64()*stddev+mean)
	}
	return res
}

func percentile(pct float64, samples []float64) float64 {
	s := sort.Float64Slice(samples)
	s.Sort()
	return s[int(float64(len(s))*pct)]
}

func Test99Percentile(t *testing.T) {
	z99 := 2.33
	e := 0.1
	p := percentile(0.99, normDistSlice(1000000, 1.0, 0.0))

	if p+e < z99 || p-e > z99 {
		t.Fatalf("99th percentile (%f) of normal distribution doesn't match the z-score of %f within an error %f", p, z99, e)
	}
}

type item struct {
	v           float64
	rank, delta int
}

type quantile struct {
	q float64 // phi like 0.50 (mean) or 0.99th
	e float64 // bias like 0.1 or 0.001
}

type quantileEstimator struct {
	samples      *list.List // <rank>
	targets      []quantile
	buffer       []float64
	observations int
}

func newBiasedQuantileEstimator(targets ...quantile) *quantileEstimator {
	return &quantileEstimator{
		targets: targets,
		samples: list.New(),
		buffer:  make([]float64, 0, 500),
	}
}

func (est *quantileEstimator) minError(rank int, n int) int {
	min := float64(n + 1)

	for _, quantile := range est.targets {
		var err float64

		if rank < int(quantile.q*float64(n)) {
			err = 2 * quantile.e * (float64(n - rank)) / (1 - quantile.q)
		} else {
			err = 2 * quantile.e * float64(rank) / quantile.q
		}

		if err < min {
			min = err
		}
	}

	return int(min)
}

// merges the batch
func (est *quantileEstimator) update(batch []float64) {
	for _, v := range batch {
		rank := 0

		e := est.samples.Front()
		if e == nil {
			e = est.samples.PushFront(&item{v, 1, 0})
			est.observations++
			continue
		}

		for e != nil {
			cur := e.Value.(*item)
			if v < cur.v {
				break
			}
			rank = rank + cur.rank
			e = e.Next()
		}

		if e == nil {
			est.samples.PushBack(&item{v, 1, est.minError(rank, est.observations) - 1})
		} else {
			est.samples.InsertBefore(&item{v, 1, est.minError(rank, est.observations) - 1}, e)
		}

		est.observations++
	}
}

func (est *quantileEstimator) flush() {
	//sort.Float64Slice(est.buffer).Sort()
	est.update(est.buffer)
	est.buffer = est.buffer[0:0]
	est.compress()
}

func (est *quantileEstimator) Update(s float64) {
	est.buffer = append(est.buffer, s)
	if len(est.buffer) == cap(est.buffer) {
		est.flush()
	}
}

func (est *quantileEstimator) compress() {
	if est.samples.Len() < 2 {
		panic("only compress filled samples")
	}

	pre := est.samples.Front()
	cur := pre.Next()

	for rank := 1; cur != nil; pre, cur = cur, cur.Next() {
		p, c := pre.Value.(*item), cur.Value.(*item)
		if p.rank+c.rank+c.delta <= est.minError(rank, est.observations) {
			c.rank += p.rank
			est.samples.Remove(pre)
		} else {
			rank += p.rank
		}

	}
}

func (est *quantileEstimator) Query(q float64) float64 {
	est.flush()

	want := int(q * float64(est.observations))
	rank := 0
	if pre := est.samples.Front(); pre != nil {
		cur := pre.Next()
		for cur != nil {
			p, c := pre.Value.(*item), cur.Value.(*item)
			rank += p.rank
			if rank+c.rank+c.delta > want+est.minError(want, est.observations)/2 {
				return p.v
			}
			pre, cur = cur, cur.Next()
		}
		return est.samples.Back().Value.(*item).v
	}
	return 0
}

func BenchmarkQuantileEstimator(b *testing.B) {
	est := newBiasedQuantileEstimator(quantile{0.01, 0.01}, quantile{0.05, 0.01}, quantile{0.50, 0.01}, quantile{0.99, 0.001})
	b.SetBytes(int64(8 * b.N + 2))
	for i := 0; i < b.N + 2; i++ {
		est.Update(rand.NormFloat64()*1.0+0.0)
	}

	println(est.samples.Len())
	println(est.Query(0.01))
	println(est.Query(0.50))
	println(est.Query(0.99))
}
<|endoftext|>


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

package py

// #include "utils.h"
import "C"

import "unsafe"

type Unicode struct {
	AbstractObject
	o C.PyUnicodeObject
}

// UnicodeType is the Type object that represents the Unicode type.
var UnicodeType = (*Type)(unsafe.Pointer(&C.PyUnicode_Type))

func unicodeCheck(obj Object) bool {
	return C.unicodeCheck(c(obj)) != 0
}

func newUnicode(obj *C.PyObject) *Unicode {
	return (*Unicode)(unsafe.Pointer(obj))
}

func NewUnicode(s string) (*Unicode, error) {
	cs := C.CString(s)
	defer C.free(unsafe.Pointer(cs))
	ret := C.PyUnicode_FromString(cs)
	if ret == nil {
		return nil, exception()
	}
	return newUnicode(ret), nil
}

func (u *Unicode) String() string {
	if u == nil {
		return "<nil>"
	}
	s := C.PyUnicode_AsUTF8String(c(u))
	if s == nil {
		panic(exception())
	}
	ret := C.PyString_AsString(s)
	if ret == nil {
		panic(exception())
	}
	return C.GoString(ret)
}

func (u *Unicode) Size() int64 {
	ret := C.PyUnicode_GetSize(c(u))
	return int64(ret)
}

func (u *Unicode) EncodeString(encoding, errors string) (Object, error) {
	var cEncoding, cErrors *C.char
	if encoding == "" {
		cEncoding = C.CString(encoding)
		defer C.free(unsafe.Pointer(cEncoding))
	}
	if errors != "" {
		cErrors = C.CString(errors)
		defer C.free(unsafe.Pointer(cErrors))
	}
	ret := C.PyUnicode_AsEncodedString(c(u), cEncoding, cErrors)
	return obj2ObjErr(ret)
}

func (u *Unicode) UTF8String() (Object, error) {
	ret := C.PyUnicode_AsUTF8String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) UTF16String() (Object, error) {
	ret := C.PyUnicode_AsUTF16String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) UTF32String() (Object, error) {
	ret := C.PyUnicode_AsUTF32String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) UnicodeEscapeString() (Object, error) {
	ret := C.PyUnicode_AsUnicodeEscapeString(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) RawUnicodeEscapeString() (Object, error) {
	ret := C.PyUnicode_AsRawUnicodeEscapeString(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) Latin1String() (Object, error) {
	ret := C.PyUnicode_AsLatin1String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) ASCIIString() (Object, error) {
	ret := C.PyUnicode_AsASCIIString(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) CharmapString(mapping Object) (Object, error) {
	ret := C.PyUnicode_AsCharmapString(c(u), c(mapping))
	return obj2ObjErr(ret)
}

func (u *Unicode) Encode(encoding, errors string) (Object, error) {
	var cEncoding, cErrors *C.char
	if encoding == "" {
		cEncoding = C.CString(encoding)
		defer C.free(unsafe.Pointer(cEncoding))
	}
	if errors != "" {
		cErrors = C.CString(errors)
		defer C.free(unsafe.Pointer(cErrors))
	}
	ret := C.PyUnicode_AsEncodedObject(c(u), cEncoding, cErrors)
	return obj2ObjErr(ret)
}

func (left *Unicode) Concat(right Object) (Object, error) {
	ret := C.PyUnicode_Concat(c(left), c(right))
	return obj2ObjErr(ret)
}

func (u *Unicode) Split(sep Object, maxsplit int64) (Object, error) {
	ret := C.PyUnicode_Split(c(u), c(sep), C.Py_ssize_t(maxsplit))
	return obj2ObjErr(ret)
}

func (u *Unicode) Splitlines(keepend bool) (Object, error) {
	cKeepend := C.int(0)
	if keepend {
		cKeepend = 1
	}
	ret := C.PyUnicode_Splitlines(c(u), cKeepend)
	return obj2ObjErr(ret)
}

func (u *Unicode) Translate(table Object, errors string) (Object, error) {
	var cErrors *C.char
	if errors != "" {
		cErrors = C.CString(errors)
		defer C.free(unsafe.Pointer(cErrors))
	}
	ret := C.PyUnicode_Translate(c(u), c(table), cErrors)
	return obj2ObjErr(ret)
}

func (u *Unicode) Join(seq Object) (Object, error) {
	ret := C.PyUnicode_Join(c(u), c(seq))
	return obj2ObjErr(ret)
}

func (u *Unicode) Tailmatch(substr Object, start, end int64, direction int) (bool, error) {
	ret := C.PyUnicode_Tailmatch(c(u), c(substr), C.Py_ssize_t(start), C.Py_ssize_t(end), C.int(direction))
	return int2BoolErr(C.int(ret))
}

func (u *Unicode) Find(substr Object, start, end int64, direction int) (int64, bool, error) {
	ret := C.PyUnicode_Find(c(u), c(substr), C.Py_ssize_t(start), C.Py_ssize_t(end), C.int(direction))
	if ret >= 0 {
		return int64(ret), true, nil
	} else if ret == -1 {
		return 0, false, nil
	}
	return 0, false, exception()
}

func (u *Unicode) Count(substr Object, start, end int64) (int64, error) {
	ret := C.PyUnicode_Count(c(u), c(substr), C.Py_ssize_t(start), C.Py_ssize_t(end))
	return ssize_t2Int64Err(ret)
}

func (u *Unicode) Replace(substr, replstr Object, maxcount int64) (Object, error) {
	ret := C.PyUnicode_Replace(c(u), c(substr), c(replstr), C.Py_ssize_t(maxcount))
	return obj2ObjErr(ret)
}

func (u *Unicode) Compare(right Object) (int, error) {
	ret := C.PyUnicode_Compare(c(u), c(right))
	return int(ret), exception()
}

func (u *Unicode) RichCompare(right Object, op Op) (Object, error) {
	ret := C.PyUnicode_RichCompare(c(u), c(right), C.int(op))
	return obj2ObjErr(ret)
}

func (u *Unicode) Format(args *Tuple) (*Unicode, error) {
	ret := C.PyUnicode_Format(c(u), c(args))
	if ret == nil {
		return nil, exception()
	}
	return newUnicode(ret), nil
}

func (u *Unicode) Contains(element Object) (bool, error) {
	ret := C.PyUnicode_Contains(c(u), c(element))
	return int2BoolErr(ret)
}
<commit_msg>Add missing decref<commit_after>// Copyright 2011 Julian Phillips.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package py

// #include "utils.h"
import "C"

import "unsafe"

type Unicode struct {
	AbstractObject
	o C.PyUnicodeObject
}

// UnicodeType is the Type object that represents the Unicode type.
var UnicodeType = (*Type)(unsafe.Pointer(&C.PyUnicode_Type))

func unicodeCheck(obj Object) bool {
	return C.unicodeCheck(c(obj)) != 0
}

func newUnicode(obj *C.PyObject) *Unicode {
	return (*Unicode)(unsafe.Pointer(obj))
}

func NewUnicode(s string) (*Unicode, error) {
	cs := C.CString(s)
	defer C.free(unsafe.Pointer(cs))
	ret := C.PyUnicode_FromString(cs)
	if ret == nil {
		return nil, exception()
	}
	return newUnicode(ret), nil
}

func (u *Unicode) String() string {
	if u == nil {
		return "<nil>"
	}
	s := C.PyUnicode_AsUTF8String(c(u))
	if s == nil {
		panic(exception())
	}
	defer C.decref(s)
	ret := C.PyString_AsString(s)
	if ret == nil {
		panic(exception())
	}
	return C.GoString(ret)
}

func (u *Unicode) Size() int64 {
	ret := C.PyUnicode_GetSize(c(u))
	return int64(ret)
}

func (u *Unicode) EncodeString(encoding, errors string) (Object, error) {
	var cEncoding, cErrors *C.char
	if encoding == "" {
		cEncoding = C.CString(encoding)
		defer C.free(unsafe.Pointer(cEncoding))
	}
	if errors != "" {
		cErrors = C.CString(errors)
		defer C.free(unsafe.Pointer(cErrors))
	}
	ret := C.PyUnicode_AsEncodedString(c(u), cEncoding, cErrors)
	return obj2ObjErr(ret)
}

func (u *Unicode) UTF8String() (Object, error) {
	ret := C.PyUnicode_AsUTF8String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) UTF16String() (Object, error) {
	ret := C.PyUnicode_AsUTF16String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) UTF32String() (Object, error) {
	ret := C.PyUnicode_AsUTF32String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) UnicodeEscapeString() (Object, error) {
	ret := C.PyUnicode_AsUnicodeEscapeString(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) RawUnicodeEscapeString() (Object, error) {
	ret := C.PyUnicode_AsRawUnicodeEscapeString(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) Latin1String() (Object, error) {
	ret := C.PyUnicode_AsLatin1String(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) ASCIIString() (Object, error) {
	ret := C.PyUnicode_AsASCIIString(c(u))
	return obj2ObjErr(ret)
}

func (u *Unicode) CharmapString(mapping Object) (Object, error) {
	ret := C.PyUnicode_AsCharmapString(c(u), c(mapping))
	return obj2ObjErr(ret)
}

func (u *Unicode) Encode(encoding, errors string) (Object, error) {
	var cEncoding, cErrors *C.char
	if encoding == "" {
		cEncoding = C.CString(encoding)
		defer C.free(unsafe.Pointer(cEncoding))
	}
	if errors != "" {
		cErrors = C.CString(errors)
		defer C.free(unsafe.Pointer(cErrors))
	}
	ret := C.PyUnicode_AsEncodedObject(c(u), cEncoding, cErrors)
	return obj2ObjErr(ret)
}

func (left *Unicode) Concat(right Object) (Object, error) {
	ret := C.PyUnicode_Concat(c(left), c(right))
	return obj2ObjErr(ret)
}

func (u *Unicode) Split(sep Object, maxsplit int64) (Object, error) {
	ret := C.PyUnicode_Split(c(u), c(sep), C.Py_ssize_t(maxsplit))
	return obj2ObjErr(ret)
}

func (u *Unicode) Splitlines(keepend bool) (Object, error) {
	cKeepend := C.int(0)
	if keepend {
		cKeepend = 1
	}
	ret := C.PyUnicode_Splitlines(c(u), cKeepend)
	return obj2ObjErr(ret)
}

func (u *Unicode) Translate(table Object, errors string) (Object, error) {
	var cErrors *C.char
	if errors != "" {
		cErrors = C.CString(errors)
		defer C.free(unsafe.Pointer(cErrors))
	}
	ret := C.PyUnicode_Translate(c(u), c(table), cErrors)
	return obj2ObjErr(ret)
}

func (u *Unicode) Join(seq Object) (Object, error) {
	ret := C.PyUnicode_Join(c(u), c(seq))
	return obj2ObjErr(ret)
}

func (u *Unicode) Tailmatch(substr Object, start, end int64, direction int) (bool, error) {
	ret := C.PyUnicode_Tailmatch(c(u), c(substr), C.Py_ssize_t(start), C.Py_ssize_t(end), C.int(direction))
	return int2BoolErr(C.int(ret))
}

func (u *Unicode) Find(substr Object, start, end int64, direction int) (int64, bool, error) {
	ret := C.PyUnicode_Find(c(u), c(substr), C.Py_ssize_t(start), C.Py_ssize_t(end), C.int(direction))
	if ret >= 0 {
		return int64(ret), true, nil
	} else if ret == -1 {
		return 0, false, nil
	}
	return 0, false, exception()
}

func (u *Unicode) Count(substr Object, start, end int64) (int64, error) {
	ret := C.PyUnicode_Count(c(u), c(substr), C.Py_ssize_t(start), C.Py_ssize_t(end))
	return ssize_t2Int64Err(ret)
}

func (u *Unicode) Replace(substr, replstr Object, maxcount int64) (Object, error) {
	ret := C.PyUnicode_Replace(c(u), c(substr), c(replstr), C.Py_ssize_t(maxcount))
	return obj2ObjErr(ret)
}

func (u *Unicode) Compare(right Object) (int, error) {
	ret := C.PyUnicode_Compare(c(u), c(right))
	return int(ret), exception()
}

func (u *Unicode) RichCompare(right Object, op Op) (Object, error) {
	ret := C.PyUnicode_RichCompare(c(u), c(right), C.int(op))
	return obj2ObjErr(ret)
}

func (u *Unicode) Format(args *Tuple) (*Unicode, error) {
	ret := C.PyUnicode_Format(c(u), c(args))
	if ret == nil {
		return nil, exception()
	}
	return newUnicode(ret), nil
}

func (u *Unicode) Contains(element Object) (bool, error) {
	ret := C.PyUnicode_Contains(c(u), c(element))
	return int2BoolErr(ret)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"time"
)

import "stathat.com/c/consistent"

// BUFFERSIZE controls the size of the [...]byte array used to read UDP data
// off the wire and into local memory.  Metrics are separated by \n
// characters.  This buffer is passed to a handler to proxy out the metrics
// to the real statsd daemons.
const BUFFERSIZE int = 1 * 1024 * 1024 // 1MiB

// packetLen is the size in bytes of data we stuff into one packet before
// sending it to statsd.  This must be lower than the MTU, IPv4 header size
// and UDP header size.
const packetLen int = 1000

// prefix is the string that will be prefixed onto self generated stats.
// Such as <prefix>.statsProcessed.  Default is "statsrelay"
var prefix string

// udpAddr is a mapping of HOST:PORT:INSTANCE to a UDPAddr object
var udpAddr = make(map[string]*net.UDPAddr)

// hashRing is our consistent hashing ring.
var hashRing *consistent.Consistent

// totalMetrics tracks the totall number of metrics processed
var totalMetrics int = 0

// totalMetricsLock is a mutex gaurding totalMetrics
var totalMetricsLock sync.Mutex

// Time we began
var epochTime int64

// Verbose/Debug output
var verbose bool

// sockBufferMaxSize() returns the maximum size that the UDP receive buffer
// in the kernel can be set to.  In bytes.
func sockBufferMaxSize() int {

	// XXX: This is Linux-only most likely
	data, err := ioutil.ReadFile("/proc/sys/net/core/rmem_max")
	if err != nil {
		log.Panicln(err)
	}

	data = bytes.TrimRight(data, "\n\r")
	i, err := strconv.Atoi(string(data))
	if err != nil {
		log.Printf("Could not parse /proc/sys/net/core/rmem_max\n")
		log.Fatalln(err)
	}

	return i
}

// getMetricName() parses the given []byte metric as a string, extracts
// the metric key name and returns it as a string.
func getMetricName(metric []byte) (string, error) {
	// statsd metrics are of the form:
	//    KEY:VALUE|TYPE|RATE
	length := bytes.IndexByte(metric, byte(':'))
	if length == -1 {
		return "error", errors.New("Length of -1, must be invalid StatsD data")
	}
	return string(metric[:length]), nil
}

// sendPacket takes a []byte and writes that directly to a UDP socket
// that was assigned for target.
func sendPacket(buff []byte, target string) {
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		log.Panicln(err)
	}

	conn.WriteToUDP(buff, udpAddr[target])
	conn.Close()
}

// buildPacketMap() is a helper function to initiallize a map that represents
// a UDP packet currently being built for each destination we proxy to.  As
// Go forbids taking the address of an object in a map or array so the
// bytes.Buffer object must be stored in the map as a pointer rather than
// a direct object in order to call the pointer methods on it.
func buildPacketMap() map[string]*bytes.Buffer {
	members := hashRing.Members()
	hash := make(map[string]*bytes.Buffer, len(members))

	for _, v := range members {
		hash[v] = new(bytes.Buffer)
	}

	return hash
}

// handleBuff() sorts through a full buffer of metrics and batches metrics
// to remote statsd daemons using a consistent hash.
func handleBuff(buff []byte) {
	packets := buildPacketMap()
	sep := []byte("\n")
	numMetrics := 0
	statsMetric := prefix + ".statsProcessed"

	for offset := 0; offset < len(buff); {
	loop:
		for offset < len(buff) {
			// Find our next value
			switch buff[offset] {
			case '\n':
				offset++
			case '\r':
				offset++
			case 0:
				offset++
			default:
				break loop
			}
		}

		size := bytes.IndexByte(buff[offset:], '\n')

		if size == -1 {
			// last metric in buffer
			size = len(buff) - offset
		}
		if size == 0 {
			// no more metrics
			break
		}

		//Check to ensure we get a metric, and not an invalid Byte sequence
		metric, err := getMetricName(buff[offset : offset+size])

		if err == nil {

			target, err := hashRing.Get(metric)
			if err != nil {
				log.Panicln(err)
			}

			// check built packet size and send if metric doesn't fit
			if packets[target].Len()+size > packetLen {
				sendPacket(packets[target].Bytes(), target)
				packets[target].Reset()
			}
			// add to packet
			packets[target].Write(buff[offset : offset+size])
			packets[target].Write(sep)

			numMetrics++
		}

		offset = offset + size + 1
	}

	if numMetrics == 0 {
		// if we haven't handled any metrics, then don't update counters/stats
		// or send packets
		return
	}

	// Update interal counter
	totalMetricsLock.Lock()
	totalMetrics = totalMetrics + numMetrics
	totalMetricsLock.Unlock()

	// Handle reporting our own stats
	stats := fmt.Sprintf("%s:%d|c\n", statsMetric, numMetrics)
	target, err := hashRing.Get(statsMetric)
	if err != nil {
		log.Panicln(err)
	}
	if packets[target].Len()+len(stats) > packetLen {
		sendPacket(packets[target].Bytes(), target)
		packets[target].Reset()
	}
	packets[target].Write([]byte(stats))

	// Empty out any remaining data
	for _, target := range hashRing.Members() {
		if packets[target].Len() > 0 {
			sendPacket(packets[target].Bytes(), target)
		}
	}

	if verbose {
		log.Printf("Procssed %d metrics. Running total: %d. Metrics/sec: %d\n",
			numMetrics, totalMetrics,
			int64(totalMetrics)/(time.Now().Unix()-epochTime))
	}
}

// readUDP() a goroutine that just reads data off of a UDP socket and fills
// buffers.  Once a buffer is full, it passes it to handleBuff().
func readUDP(ip string, port int, c chan []byte) {
	var buff *[BUFFERSIZE]byte
	var offset int
	var timeout bool
	var addr = net.UDPAddr{
		Port: port,
		IP:   net.ParseIP(ip),
	}

	log.Printf("Listening on %s:%d\n", ip, port)
	sock, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Printf("Error opening UDP socket.\n")
		log.Fatalln(err)
	}
	defer sock.Close()

	log.Printf("Setting socket read buffer size to: %d\n", sockBufferMaxSize())
	err = sock.SetReadBuffer(sockBufferMaxSize())
	if err != nil {
		log.Printf("Unable to set read buffer size on socket.  Non-fatal.")
		log.Println(err)
	}
	err = sock.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		log.Printf("Unable to set timeout on socket.\n")
		log.Fatalln(err)
	}

	if verbose {
		log.Printf("Rock and Roll!\n")
	}
	for {
		if buff == nil {
			buff = new([BUFFERSIZE]byte)
			offset = 0
			timeout = false
		}

		i, err := sock.Read(buff[offset:])
		if err == nil {
			buff[offset+i] = '\n'
			offset = offset + i + 1
		} else if err.(net.Error).Timeout() {
			timeout = true
			err = sock.SetDeadline(time.Now().Add(time.Second))
			if err != nil {
				log.Panicln(err)
			}
		} else {
			log.Printf("Read Error: %s\n", err)
			continue
		}

		if offset > BUFFERSIZE-4096 || timeout {
			// Approching make buff size
			// we use a 4KiB margin
			c <- buff[:offset]
			buff = nil
		}
	}
}

// runServer() runs and manages this daemon, deals with OS signals, and handles
// communication channels.
func runServer(host string, port int) {
	var c chan []byte = make(chan []byte, 256)
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	var sig chan os.Signal = make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	// read incming UDP packets
	go readUDP(host, port, c)

	for {
		select {
		case buff := <-c:
			//fmt.Printf("Handling %d length buffer...\n", len(buff))
			go handleBuff(buff)
		case <-sig:
			log.Printf("Signal received.  Shutting down...\n")
			log.Printf("Received %d metrics.\n", totalMetrics)
			return
		}
	}
}

func main() {
	var bindAddress string
	var port int

	flag.IntVar(&port, "port", 9125, "Port to listen on")
	flag.IntVar(&port, "p", 9125, "Port to listen on")

	flag.StringVar(&bindAddress, "bind", "0.0.0.0", "IP Address to listen on")
	flag.StringVar(&bindAddress, "b", "0.0.0.0", "IP Address to listen on")

	flag.StringVar(&prefix, "prefix", "statsrelay", "The prefix to use with self generated stats")

	flag.BoolVar(&verbose, "verbose", false, "Verbose output")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatalf("One or more host specifications are needed to locate statsd daemons.\n")
	}

	hashRing = consistent.New()
	hashRing.NumberOfReplicas = 1

	for _, v := range flag.Args() {
		var addr *net.UDPAddr
		var err error
		host := strings.Split(v, ":")

		switch len(host) {
		case 1:
			log.Printf("Invalid statsd location: %s\n", v)
			log.Fatalf("Must be of the form HOST:PORT or HOST:PORT:INSTANCE\n")
		case 2:
			addr, err = net.ResolveUDPAddr("udp", v)
			if err != nil {
				log.Printf("Error parsing HOST:PORT \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		case 3:
			addr, err = net.ResolveUDPAddr("udp", host[0]+":"+host[1])
			if err != nil {
				log.Printf("Error parsing HOST:PORT:INSTANCE \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		default:
			log.Fatalf("Unrecongnized host specification: %s\n", v)
		}

		if addr != nil {
			udpAddr[v] = addr
			hashRing.Add(v)
		}
	}

	epochTime = time.Now().Unix()
	runServer(bindAddress, port)

	log.Printf("Normal shutdown.\n")
}
<commit_msg>Verbose now logs where metrics are routed.<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"time"
)

import "stathat.com/c/consistent"

// BUFFERSIZE controls the size of the [...]byte array used to read UDP data
// off the wire and into local memory.  Metrics are separated by \n
// characters.  This buffer is passed to a handler to proxy out the metrics
// to the real statsd daemons.
const BUFFERSIZE int = 1 * 1024 * 1024 // 1MiB

// packetLen is the size in bytes of data we stuff into one packet before
// sending it to statsd.  This must be lower than the MTU, IPv4 header size
// and UDP header size.
const packetLen int = 1000

// prefix is the string that will be prefixed onto self generated stats.
// Such as <prefix>.statsProcessed.  Default is "statsrelay"
var prefix string

// udpAddr is a mapping of HOST:PORT:INSTANCE to a UDPAddr object
var udpAddr = make(map[string]*net.UDPAddr)

// hashRing is our consistent hashing ring.
var hashRing *consistent.Consistent

// totalMetrics tracks the totall number of metrics processed
var totalMetrics int = 0

// totalMetricsLock is a mutex gaurding totalMetrics
var totalMetricsLock sync.Mutex

// Time we began
var epochTime int64

// Verbose/Debug output
var verbose bool

// sockBufferMaxSize() returns the maximum size that the UDP receive buffer
// in the kernel can be set to.  In bytes.
func sockBufferMaxSize() int {

	// XXX: This is Linux-only most likely
	data, err := ioutil.ReadFile("/proc/sys/net/core/rmem_max")
	if err != nil {
		log.Panicln(err)
	}

	data = bytes.TrimRight(data, "\n\r")
	i, err := strconv.Atoi(string(data))
	if err != nil {
		log.Printf("Could not parse /proc/sys/net/core/rmem_max\n")
		log.Fatalln(err)
	}

	return i
}

// getMetricName() parses the given []byte metric as a string, extracts
// the metric key name and returns it as a string.
func getMetricName(metric []byte) (string, error) {
	// statsd metrics are of the form:
	//    KEY:VALUE|TYPE|RATE
	length := bytes.IndexByte(metric, byte(':'))
	if length == -1 {
		return "error", errors.New("Length of -1, must be invalid StatsD data")
	}
	return string(metric[:length]), nil
}

// sendPacket takes a []byte and writes that directly to a UDP socket
// that was assigned for target.
func sendPacket(buff []byte, target string) {
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		log.Panicln(err)
	}

	conn.WriteToUDP(buff, udpAddr[target])
	conn.Close()
}

// buildPacketMap() is a helper function to initiallize a map that represents
// a UDP packet currently being built for each destination we proxy to.  As
// Go forbids taking the address of an object in a map or array so the
// bytes.Buffer object must be stored in the map as a pointer rather than
// a direct object in order to call the pointer methods on it.
func buildPacketMap() map[string]*bytes.Buffer {
	members := hashRing.Members()
	hash := make(map[string]*bytes.Buffer, len(members))

	for _, v := range members {
		hash[v] = new(bytes.Buffer)
	}

	return hash
}

// handleBuff() sorts through a full buffer of metrics and batches metrics
// to remote statsd daemons using a consistent hash.
func handleBuff(buff []byte) {
	packets := buildPacketMap()
	sep := []byte("\n")
	numMetrics := 0
	statsMetric := prefix + ".statsProcessed"

	for offset := 0; offset < len(buff); {
	loop:
		for offset < len(buff) {
			// Find our next value
			switch buff[offset] {
			case '\n':
				offset++
			case '\r':
				offset++
			case 0:
				offset++
			default:
				break loop
			}
		}

		size := bytes.IndexByte(buff[offset:], '\n')

		if size == -1 {
			// last metric in buffer
			size = len(buff) - offset
		}
		if size == 0 {
			// no more metrics
			break
		}

		//Check to ensure we get a metric, and not an invalid Byte sequence
		metric, err := getMetricName(buff[offset : offset+size])

		if err == nil {

			target, err := hashRing.Get(metric)
			if err != nil {
				log.Panicln(err)
			}
			if verbose {
				log.Printf("Sending %s to %s", metric, target)
			}

			// check built packet size and send if metric doesn't fit
			if packets[target].Len()+size > packetLen {
				sendPacket(packets[target].Bytes(), target)
				packets[target].Reset()
			}
			// add to packet
			packets[target].Write(buff[offset : offset+size])
			packets[target].Write(sep)

			numMetrics++
		}

		offset = offset + size + 1
	}

	if numMetrics == 0 {
		// if we haven't handled any metrics, then don't update counters/stats
		// or send packets
		return
	}

	// Update interal counter
	totalMetricsLock.Lock()
	totalMetrics = totalMetrics + numMetrics
	totalMetricsLock.Unlock()

	// Handle reporting our own stats
	stats := fmt.Sprintf("%s:%d|c\n", statsMetric, numMetrics)
	target, err := hashRing.Get(statsMetric)
	if err != nil {
		log.Panicln(err)
	}
	if packets[target].Len()+len(stats) > packetLen {
		sendPacket(packets[target].Bytes(), target)
		packets[target].Reset()
	}
	packets[target].Write([]byte(stats))

	// Empty out any remaining data
	for _, target := range hashRing.Members() {
		if packets[target].Len() > 0 {
			sendPacket(packets[target].Bytes(), target)
		}
	}

	if verbose {
		log.Printf("Procssed %d metrics. Running total: %d. Metrics/sec: %d\n",
			numMetrics, totalMetrics,
			int64(totalMetrics)/(time.Now().Unix()-epochTime))
	}
}

// readUDP() a goroutine that just reads data off of a UDP socket and fills
// buffers.  Once a buffer is full, it passes it to handleBuff().
func readUDP(ip string, port int, c chan []byte) {
	var buff *[BUFFERSIZE]byte
	var offset int
	var timeout bool
	var addr = net.UDPAddr{
		Port: port,
		IP:   net.ParseIP(ip),
	}

	log.Printf("Listening on %s:%d\n", ip, port)
	sock, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Printf("Error opening UDP socket.\n")
		log.Fatalln(err)
	}
	defer sock.Close()

	log.Printf("Setting socket read buffer size to: %d\n", sockBufferMaxSize())
	err = sock.SetReadBuffer(sockBufferMaxSize())
	if err != nil {
		log.Printf("Unable to set read buffer size on socket.  Non-fatal.")
		log.Println(err)
	}
	err = sock.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		log.Printf("Unable to set timeout on socket.\n")
		log.Fatalln(err)
	}

	if verbose {
		log.Printf("Rock and Roll!\n")
	}
	for {
		if buff == nil {
			buff = new([BUFFERSIZE]byte)
			offset = 0
			timeout = false
		}

		i, err := sock.Read(buff[offset:])
		if err == nil {
			buff[offset+i] = '\n'
			offset = offset + i + 1
		} else if err.(net.Error).Timeout() {
			timeout = true
			err = sock.SetDeadline(time.Now().Add(time.Second))
			if err != nil {
				log.Panicln(err)
			}
		} else {
			log.Printf("Read Error: %s\n", err)
			continue
		}

		if offset > BUFFERSIZE-4096 || timeout {
			// Approching make buff size
			// we use a 4KiB margin
			c <- buff[:offset]
			buff = nil
		}
	}
}

// runServer() runs and manages this daemon, deals with OS signals, and handles
// communication channels.
func runServer(host string, port int) {
	var c chan []byte = make(chan []byte, 256)
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	var sig chan os.Signal = make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	// read incming UDP packets
	go readUDP(host, port, c)

	for {
		select {
		case buff := <-c:
			//fmt.Printf("Handling %d length buffer...\n", len(buff))
			go handleBuff(buff)
		case <-sig:
			log.Printf("Signal received.  Shutting down...\n")
			log.Printf("Received %d metrics.\n", totalMetrics)
			return
		}
	}
}

func main() {
	var bindAddress string
	var port int

	flag.IntVar(&port, "port", 9125, "Port to listen on")
	flag.IntVar(&port, "p", 9125, "Port to listen on")

	flag.StringVar(&bindAddress, "bind", "0.0.0.0", "IP Address to listen on")
	flag.StringVar(&bindAddress, "b", "0.0.0.0", "IP Address to listen on")

	flag.StringVar(&prefix, "prefix", "statsrelay", "The prefix to use with self generated stats")

	flag.BoolVar(&verbose, "verbose", false, "Verbose output")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatalf("One or more host specifications are needed to locate statsd daemons.\n")
	}

	hashRing = consistent.New()
	hashRing.NumberOfReplicas = 1

	for _, v := range flag.Args() {
		var addr *net.UDPAddr
		var err error
		host := strings.Split(v, ":")

		switch len(host) {
		case 1:
			log.Printf("Invalid statsd location: %s\n", v)
			log.Fatalf("Must be of the form HOST:PORT or HOST:PORT:INSTANCE\n")
		case 2:
			addr, err = net.ResolveUDPAddr("udp", v)
			if err != nil {
				log.Printf("Error parsing HOST:PORT \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		case 3:
			addr, err = net.ResolveUDPAddr("udp", host[0]+":"+host[1])
			if err != nil {
				log.Printf("Error parsing HOST:PORT:INSTANCE \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		default:
			log.Fatalf("Unrecongnized host specification: %s\n", v)
		}

		if addr != nil {
			udpAddr[v] = addr
			hashRing.Add(v)
		}
	}

	epochTime = time.Now().Unix()
	runServer(bindAddress, port)

	log.Printf("Normal shutdown.\n")
}
<|endoftext|>