<commit_before>package bubbles

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	// DefaultMaxDocumentsPerBatch is the number of documents a batch needs to
	// have before it is send. This is per connection.
	DefaultMaxDocumentsPerBatch = 1000

	// DefaultFlushTimeout is the maximum time we batch something before we try
	// to send it to a server.
	DefaultFlushTimeout = 10 * time.Second

	// DefaultServerTimeout is the time we give ElasticSearch to respond.
	DefaultServerTimeout = 1 * time.Minute

	// DefaultConnCount is the number of connections per hosts.
	DefaultConnCount = 2

	// backoffTimeoutRatio determines when we start backing off.
	backoffTimeoutRatio = 2

	serverErrorWait    = 500 * time.Millisecond
	serverErrorWaitMax = 16 * time.Second

	defaultElasticSearchPort = "9200"
)

var (
	errInvalidResponse = errors.New("invalid response")
)

// Bubbles is the main struct to control a queue of Actions going to the
// ElasticSearch servers.
type Bubbles struct {
	q                chan Action
	retryQ           chan Action
	quit             chan struct{}
	wg               sync.WaitGroup
	maxDocumentCount int
	connCount        int
	flushTimeout     time.Duration
	serverTimeout    time.Duration
	c                Counter
	e                Errer
}

// Opt is any option to New().
type Opt func(*Bubbles)

// OptConnCount is an option to New() to specify the number of connections per
// host. The default is DefaultConnCount.
func OptConnCount(n int) Opt {
	return func(b *Bubbles) {
		b.connCount = n
	}
}

// OptFlush is an option to New() to specify the flush timeout of a batch. The
// default is DefaultFlushTimeout.
func OptFlush(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.flushTimeout = d
	}
}

// OptServerTimeout is an option to New() to specify the timeout of a single
// batch POST to ElasticSearch. This value is also the maximum time Stop() will
// take. All actions in a bulk which is timed out will be retried. The default
// is DefaultServerTimeout.
func OptServerTimeout(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.serverTimeout = d
	}
}

// OptMaxDocs is an option to New() to specify maximum number of documents in a
// single batch. The default is DefaultMaxDocumentsPerBatch.
func OptMaxDocs(n int) Opt {
	return func(b *Bubbles) {
		b.maxDocumentCount = n
	}
}

// OptCounter is an option to New() to specify something that counts documents.
func OptCounter(c Counter) Opt {
	return func(b *Bubbles) {
		b.c = c
	}
}

// OptErrer is an option to New() to specify an error handler. The default
// handler uses the log module.
func OptErrer(e Errer) Opt {
	return func(b *Bubbles) {
		b.e = e
	}
}

// New makes a new ElasticSearch bulk inserter. It needs a list with 'ip' or
// 'ip:port' addresses, options are added via the Opt* functions.
func New(addrs []string, opts ...Opt) *Bubbles {
	b := Bubbles{
		q:                make(chan Action),
		quit:             make(chan struct{}),
		maxDocumentCount: DefaultMaxDocumentsPerBatch,
		connCount:        DefaultConnCount,
		flushTimeout:     DefaultFlushTimeout,
		serverTimeout:    DefaultServerTimeout,
		c:                DefaultCounter{},
		e:                DefaultErrer{},
	}
	for _, o := range opts {
		o(&b)
	}
	b.retryQ = make(chan Action, len(addrs)*b.connCount*b.maxDocumentCount)

	cl := &http.Client{
		Timeout: b.serverTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return errors.New("no redirect")
		},
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   b.serverTimeout,
				KeepAlive: 30 * time.Second,
			}).Dial,
			MaxIdleConnsPerHost: b.connCount,
			DisableCompression:  false,
		},
	}
	// Start a go routine per connection per host
	for _, a := range addrs {
		addr := withPort(a, defaultElasticSearchPort)
		for i := 0; i < b.connCount; i++ {
			b.wg.Add(1)
			go func(a string) {
				client(&b, cl, a)
				b.wg.Done()
			}(addr)
		}
	}
	return &b
}

// Enqueue returns the queue to add Actions in a routine. It will block if all bulk
// processors are busy.
func (b *Bubbles) Enqueue() chan<- Action {
	return b.q
}

// Stop shuts down all ElasticSearch clients. It'll return all Action entries
// which were not yet processed, or were up for a retry. It can take
// OptServerTimeout to complete.
func (b *Bubbles) Stop() []Action {
	close(b.quit)
	// There is no explicit timeout, we rely on b.serverTimeout to shut down
	// everything.
	b.wg.Wait()

	// Collect and return elements which are in flight.
	close(b.retryQ)
	close(b.q)
	pending := make([]Action, 0, len(b.q)+len(b.retryQ))
	for a := range b.q {
		pending = append(pending, a)
	}
	for a := range b.retryQ {
		pending = append(pending, a)
	}
	return pending
}

type backoff struct {
	level        uint8
	max          uint8
	maxBatchSize int
}

func newBackoff(maxBatchSize int) *backoff {
	b := &backoff{
		maxBatchSize: maxBatchSize,
	}
	// Max out when both components do.
	for ; (b.wait() < serverErrorWaitMax || b.size() > 1) && b.level < math.MaxUint8; b.level++ {
	}
	b.max = b.level
	b.level = 0
	return b
}

// wait calculates the delay based on the current backoff level.
func (b *backoff) wait() time.Duration {
	if b.level == 0 {
		return 0 * time.Second
	}
	w := (1 << (b.level - 1)) * serverErrorWait
	if w >= serverErrorWaitMax {
		return serverErrorWaitMax
	}
	return w
}

// batchSize calculates the batchsize based on the current backoff level.
func (b *backoff) size() int {
	s := b.maxBatchSize / (1 << b.level)
	if s <= 1 {
		return 1
	}
	return s
}

// inc increases the backoff level.
func (b *backoff) inc() {
	if b.level < b.max {
		b.level++
	}
}

// dec decreases the backoff level.
func (b *backoff) dec() {
	if b.level > 0 {
		b.level--
	}
}

// client talks to ElasticSearch. This runs in a go routine in a loop and deals
// with a single ElasticSearch address.
func client(b *Bubbles, cl *http.Client, addr string) {
	url := fmt.Sprintf("http://%s/_bulk", addr)

	backoff := newBackoff(b.maxDocumentCount)
	backoffTime := b.serverTimeout / backoffTimeoutRatio
	for {
		select {
		case <-b.quit:
			return
		case <-time.After(backoff.wait()):
		}
		trouble, batchTime := runBatch(b, cl, url, backoff.size())
		if trouble || batchTime > backoffTime {
			backoff.inc()
			b.c.Timeout()
		} else {
			backoff.dec()
		}
		b.c.BatchTime(batchTime)
	}
}

// runBatch gathers and deals with a batch of actions. It returns
// whether there was trouble, and how long the actual request took.
func runBatch(b *Bubbles, cl *http.Client, url string, batchSize int) (bool, time.Duration) {
	actions := make([]Action, 0, b.maxDocumentCount)
	// First use all retry actions.
retry:
	for len(actions) < batchSize {
		select {
		case a := <-b.retryQ:
			actions = append(actions, a)
		default:
			// no more retry actions queued
			break retry
		}
	}

	var t <-chan time.Time
gather:
	for len(actions) < batchSize {
		if t == nil && len(actions) > 0 {
			// Set timeout on the first element we read
			t = time.After(b.flushTimeout)
		}
		select {
		case <-b.quit:
			for _, a := range actions {
				b.retryQ <- a
			}
			return false, 0
		case <-t:
			// this case is not enabled until we've got an action
			break gather
		case a := <-b.retryQ:
			actions = append(actions, a)
		case a := <-b.q:
			actions = append(actions, a)
		}
	}

	t0 := time.Now()
	res, err := postActions(b.c, cl, url, actions, b.quit)
	dt := time.Since(t0)
	if err != nil {
		// A server error. Retry these actions later.
		b.e.Error(err)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}

	// Server has accepted the request an sich, but there can be errors in the
	// individual actions.
	if !res.Errors {
		// Simple case, no errors present.
		return false, dt
	}

	// Invalid response from ElasticSearch.
	if len(actions) != len(res.Items) {
		b.e.Error(errInvalidResponse)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}
	// Figure out which actions have errors.
	for i, e := range res.Items {
		a := actions[i]
		el, ok := e[string(a.Type)]
		if !ok {
			// Unexpected reply from ElasticSearch.
			b.e.Error(errInvalidResponse)
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
			continue
		}

		c := el.Status
		switch {
		case c >= 200 && c < 300:
			// Document accepted by ElasticSearch.
		case c == 429 || (c >= 500 && c < 600):
			// Server error. Retry it.
			// We get a 429 when the bulk queue is full, which we just retry as
			// well.
			b.e.Warning(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("transient error %d: %s", c, el.Error),
				Server:     url,
			})
			b.c.Retry(RetryTransient, a.Type, len(a.Document))
			b.retryQ <- a
		case c >= 400 && c < 500:
			// Some error. Nothing we can do with it.
			b.e.Error(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("error %d: %s", c, el.Error),
				Server:     url,
			})
		default:
			// No idea.
			b.e.Error(fmt.Errorf("unwelcome response %d: %s", c, el.Error))
		}
	}
	return true, dt
}

type bulkRes struct {
	Took   int  `json:"took"`
	Errors bool `json:"errors"`
	Items  []map[string]struct {
		Index   string `json:"_index"`
		Type    string `json:"_type"`
		ID      string `json:"_id"`
		Version int    `json:"_version"`
		Status  int    `json:"status"`
		Error   string `json:"error"`
	} `json:"items"`
}

func interruptibleDo(cl *http.Client, req *http.Request, interrupt <-chan struct{}) (*http.Response, error) {
	done := make(chan struct{})
	go func() {
		select {
		case <-interrupt:
			cl.Transport.(*http.Transport).CancelRequest(req)
		case <-done:
		}
	}()
	defer close(done)
	return cl.Do(req)
}

func postActions(c Counter, cl *http.Client, url string, actions []Action, quit <-chan struct{}) (*bulkRes, error) {
	buf := bytes.Buffer{}
	for _, a := range actions {
		c.Send(a.Type, len(a.Document))
		buf.Write(a.Buf())
	}
	c.SendTotal(buf.Len())

	// This doesn't Chunk.
	req, err := http.NewRequest("POST", url, &buf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := interruptibleDo(cl, req, quit)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(body))
	}

	var bulk bulkRes
	if err := json.Unmarshal(body, &bulk); err != nil {
		return nil, err
	}
	return &bulk, nil
}

// ActionError wraps an Action we won't retry. It implements the error interface.
type ActionError struct {
	Action     Action
	StatusCode int
	Msg        string
	Server     string
}

func (e ActionError) Error() string {
	return fmt.Sprintf("%s: %s %s", e.Server, e.Action.Type, e.Msg)
}

// withPort adds a default port to an address string.
func withPort(a, port string) string {
	if _, _, err := net.SplitHostPort(a); err != nil {
		// no port found.
		return net.JoinHostPort(a, port)
	}
	return a
}
<commit_msg>A package documentation fragment.<commit_after>// Package bubbles is an Elasticsearch bulk insert client.
//
// It connects to an Elasticsearch cluster via the bulk API
// (http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-bulk.html).
// Actions are batched into bulk requests. Actions which resulted
// in an error are retried individually.
package bubbles

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	// DefaultMaxDocumentsPerBatch is the number of documents a batch needs to
	// have before it is send. This is per connection.
	DefaultMaxDocumentsPerBatch = 1000

	// DefaultFlushTimeout is the maximum time we batch something before we try
	// to send it to a server.
	DefaultFlushTimeout = 10 * time.Second

	// DefaultServerTimeout is the time we give ElasticSearch to respond.
	DefaultServerTimeout = 1 * time.Minute

	// DefaultConnCount is the number of connections per hosts.
	DefaultConnCount = 2

	// backoffTimeoutRatio determines when we start backing off.
	backoffTimeoutRatio = 2

	serverErrorWait    = 500 * time.Millisecond
	serverErrorWaitMax = 16 * time.Second

	defaultElasticSearchPort = "9200"
)

var (
	errInvalidResponse = errors.New("invalid response")
)

// Bubbles is the main struct to control a queue of Actions going to the
// ElasticSearch servers.
type Bubbles struct {
	q                chan Action
	retryQ           chan Action
	quit             chan struct{}
	wg               sync.WaitGroup
	maxDocumentCount int
	connCount        int
	flushTimeout     time.Duration
	serverTimeout    time.Duration
	c                Counter
	e                Errer
}

// Opt is any option to New().
type Opt func(*Bubbles)

// OptConnCount is an option to New() to specify the number of connections per
// host. The default is DefaultConnCount.
func OptConnCount(n int) Opt {
	return func(b *Bubbles) {
		b.connCount = n
	}
}

// OptFlush is an option to New() to specify the flush timeout of a batch. The
// default is DefaultFlushTimeout.
func OptFlush(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.flushTimeout = d
	}
}

// OptServerTimeout is an option to New() to specify the timeout of a single
// batch POST to ElasticSearch. This value is also the maximum time Stop() will
// take. All actions in a bulk which is timed out will be retried. The default
// is DefaultServerTimeout.
func OptServerTimeout(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.serverTimeout = d
	}
}

// OptMaxDocs is an option to New() to specify maximum number of documents in a
// single batch. The default is DefaultMaxDocumentsPerBatch.
func OptMaxDocs(n int) Opt {
	return func(b *Bubbles) {
		b.maxDocumentCount = n
	}
}

// OptCounter is an option to New() to specify something that counts documents.
func OptCounter(c Counter) Opt {
	return func(b *Bubbles) {
		b.c = c
	}
}

// OptErrer is an option to New() to specify an error handler. The default
// handler uses the log module.
func OptErrer(e Errer) Opt {
	return func(b *Bubbles) {
		b.e = e
	}
}

// New makes a new ElasticSearch bulk inserter. It needs a list with 'ip' or
// 'ip:port' addresses, options are added via the Opt* functions.
func New(addrs []string, opts ...Opt) *Bubbles {
	b := Bubbles{
		q:                make(chan Action),
		quit:             make(chan struct{}),
		maxDocumentCount: DefaultMaxDocumentsPerBatch,
		connCount:        DefaultConnCount,
		flushTimeout:     DefaultFlushTimeout,
		serverTimeout:    DefaultServerTimeout,
		c:                DefaultCounter{},
		e:                DefaultErrer{},
	}
	for _, o := range opts {
		o(&b)
	}
	b.retryQ = make(chan Action, len(addrs)*b.connCount*b.maxDocumentCount)

	cl := &http.Client{
		Timeout: b.serverTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return errors.New("no redirect")
		},
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   b.serverTimeout,
				KeepAlive: 30 * time.Second,
			}).Dial,
			MaxIdleConnsPerHost: b.connCount,
			DisableCompression:  false,
		},
	}
	// Start a go routine per connection per host
	for _, a := range addrs {
		addr := withPort(a, defaultElasticSearchPort)
		for i := 0; i < b.connCount; i++ {
			b.wg.Add(1)
			go func(a string) {
				client(&b, cl, a)
				b.wg.Done()
			}(addr)
		}
	}
	return &b
}

// Enqueue returns the queue to add Actions in a routine. It will block if all bulk
// processors are busy.
func (b *Bubbles) Enqueue() chan<- Action {
	return b.q
}

// Stop shuts down all ElasticSearch clients. It'll return all Action entries
// which were not yet processed, or were up for a retry. It can take
// OptServerTimeout to complete.
func (b *Bubbles) Stop() []Action {
	close(b.quit)
	// There is no explicit timeout, we rely on b.serverTimeout to shut down
	// everything.
	b.wg.Wait()

	// Collect and return elements which are in flight.
	close(b.retryQ)
	close(b.q)
	pending := make([]Action, 0, len(b.q)+len(b.retryQ))
	for a := range b.q {
		pending = append(pending, a)
	}
	for a := range b.retryQ {
		pending = append(pending, a)
	}
	return pending
}

type backoff struct {
	level        uint8
	max          uint8
	maxBatchSize int
}

func newBackoff(maxBatchSize int) *backoff {
	b := &backoff{
		maxBatchSize: maxBatchSize,
	}
	// Max out when both components do.
	for ; (b.wait() < serverErrorWaitMax || b.size() > 1) && b.level < math.MaxUint8; b.level++ {
	}
	b.max = b.level
	b.level = 0
	return b
}

// wait calculates the delay based on the current backoff level.
func (b *backoff) wait() time.Duration {
	if b.level == 0 {
		return 0 * time.Second
	}
	w := (1 << (b.level - 1)) * serverErrorWait
	if w >= serverErrorWaitMax {
		return serverErrorWaitMax
	}
	return w
}

// batchSize calculates the batchsize based on the current backoff level.
func (b *backoff) size() int {
	s := b.maxBatchSize / (1 << b.level)
	if s <= 1 {
		return 1
	}
	return s
}

// inc increases the backoff level.
func (b *backoff) inc() {
	if b.level < b.max {
		b.level++
	}
}

// dec decreases the backoff level.
func (b *backoff) dec() {
	if b.level > 0 {
		b.level--
	}
}

// client talks to ElasticSearch. This runs in a go routine in a loop and deals
// with a single ElasticSearch address.
func client(b *Bubbles, cl *http.Client, addr string) {
	url := fmt.Sprintf("http://%s/_bulk", addr)

	backoff := newBackoff(b.maxDocumentCount)
	backoffTime := b.serverTimeout / backoffTimeoutRatio
	for {
		select {
		case <-b.quit:
			return
		case <-time.After(backoff.wait()):
		}
		trouble, batchTime := runBatch(b, cl, url, backoff.size())
		if trouble || batchTime > backoffTime {
			backoff.inc()
			b.c.Timeout()
		} else {
			backoff.dec()
		}
		b.c.BatchTime(batchTime)
	}
}

// runBatch gathers and deals with a batch of actions. It returns
// whether there was trouble, and how long the actual request took.
func runBatch(b *Bubbles, cl *http.Client, url string, batchSize int) (bool, time.Duration) {
	actions := make([]Action, 0, b.maxDocumentCount)
	// First use all retry actions.
retry:
	for len(actions) < batchSize {
		select {
		case a := <-b.retryQ:
			actions = append(actions, a)
		default:
			// no more retry actions queued
			break retry
		}
	}

	var t <-chan time.Time
gather:
	for len(actions) < batchSize {
		if t == nil && len(actions) > 0 {
			// Set timeout on the first element we read
			t = time.After(b.flushTimeout)
		}
		select {
		case <-b.quit:
			for _, a := range actions {
				b.retryQ <- a
			}
			return false, 0
		case <-t:
			// this case is not enabled until we've got an action
			break gather
		case a := <-b.retryQ:
			actions = append(actions, a)
		case a := <-b.q:
			actions = append(actions, a)
		}
	}

	t0 := time.Now()
	res, err := postActions(b.c, cl, url, actions, b.quit)
	dt := time.Since(t0)
	if err != nil {
		// A server error. Retry these actions later.
		b.e.Error(err)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}

	// Server has accepted the request an sich, but there can be errors in the
	// individual actions.
	if !res.Errors {
		// Simple case, no errors present.
		return false, dt
	}

	// Invalid response from ElasticSearch.
	if len(actions) != len(res.Items) {
		b.e.Error(errInvalidResponse)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}
	// Figure out which actions have errors.
	for i, e := range res.Items {
		a := actions[i]
		el, ok := e[string(a.Type)]
		if !ok {
			// Unexpected reply from ElasticSearch.
			b.e.Error(errInvalidResponse)
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
			continue
		}

		c := el.Status
		switch {
		case c >= 200 && c < 300:
			// Document accepted by ElasticSearch.
		case c == 429 || (c >= 500 && c < 600):
			// Server error. Retry it.
			// We get a 429 when the bulk queue is full, which we just retry as
			// well.
			b.e.Warning(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("transient error %d: %s", c, el.Error),
				Server:     url,
			})
			b.c.Retry(RetryTransient, a.Type, len(a.Document))
			b.retryQ <- a
		case c >= 400 && c < 500:
			// Some error. Nothing we can do with it.
			b.e.Error(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("error %d: %s", c, el.Error),
				Server:     url,
			})
		default:
			// No idea.
			b.e.Error(fmt.Errorf("unwelcome response %d: %s", c, el.Error))
		}
	}
	return true, dt
}

type bulkRes struct {
	Took   int  `json:"took"`
	Errors bool `json:"errors"`
	Items  []map[string]struct {
		Index   string `json:"_index"`
		Type    string `json:"_type"`
		ID      string `json:"_id"`
		Version int    `json:"_version"`
		Status  int    `json:"status"`
		Error   string `json:"error"`
	} `json:"items"`
}

func interruptibleDo(cl *http.Client, req *http.Request, interrupt <-chan struct{}) (*http.Response, error) {
	done := make(chan struct{})
	go func() {
		select {
		case <-interrupt:
			cl.Transport.(*http.Transport).CancelRequest(req)
		case <-done:
		}
	}()
	defer close(done)
	return cl.Do(req)
}

func postActions(c Counter, cl *http.Client, url string, actions []Action, quit <-chan struct{}) (*bulkRes, error) {
	buf := bytes.Buffer{}
	for _, a := range actions {
		c.Send(a.Type, len(a.Document))
		buf.Write(a.Buf())
	}
	c.SendTotal(buf.Len())

	// This doesn't Chunk.
	req, err := http.NewRequest("POST", url, &buf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := interruptibleDo(cl, req, quit)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(body))
	}

	var bulk bulkRes
	if err := json.Unmarshal(body, &bulk); err != nil {
		return nil, err
	}
	return &bulk, nil
}

// ActionError wraps an Action we won't retry. It implements the error interface.
type ActionError struct {
	Action     Action
	StatusCode int
	Msg        string
	Server     string
}

func (e ActionError) Error() string {
	return fmt.Sprintf("%s: %s %s", e.Server, e.Action.Type, e.Msg)
}

// withPort adds a default port to an address string.
func withPort(a, port string) string {
	if _, _, err := net.SplitHostPort(a); err != nil {
		// no port found.
		return net.JoinHostPort(a, port)
	}
	return a
}
<|endoftext|>


<commit_before>//Package geozip (some description here)
package geozip

//Encode (some description here)
func Encode(latitude, longitude float64, validate bool, precision int) int64 {
	var bucket int64
	return bucket
}

//Decode (some description here)
func Decode(bucket int64) (float64, float64) {
	var latitude, longitude float64
	return latitude, longitude
}

//Validate (some description here)
func Validate(latitude, longitude float64) bool {
	var validated bool
	return validated
}
<commit_msg>Encode and validate written.<commit_after>//Package geozip (some description here)
package geozip

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/shopspring/decimal"
)

//Encode (some description here)
func Encode(latitude, longitude float64, validate bool, precision int) int64 {
	if !Validate(latitude, longitude) {
		return 0
	}
	latitudeShifted := decimal.NewFromFloat(latitude).Add(decimal.NewFromFloat(90.0))
	longitudeShifted := decimal.NewFromFloat(longitude).Add(decimal.NewFromFloat(180.0))
	latString := latitudeShifted.String()
	lonString := longitudeShifted.String()
	latString = strings.Replace(latString, ".", "", 1)
	lonString = strings.Replace(lonString, ".", "", 1)
	bucketString := zip(latString, lonString)
	bucket, err := strconv.ParseInt(bucketString, 10, 64)
	if err != nil {
		fmt.Errorf("Error parsing zipped string to int64")
	}
	//Fix precision here
	return bucket
}

//Decode (some description here)
func Decode(bucket int64) (float64, float64) {
	var latitude, longitude float64
	return latitude, longitude
}

//Validate (some description here)
func Validate(latitude, longitude float64) bool {
	if (latitude < 90.0 || latitude > -90.0) && (longitude < 180.0 || longitude > 180.0) {
		return true
	}
	return false

}

func zip(latDigits, lonDigits string) string {
	var bucketDigits string
	latDigits = resize(latDigits)
	lonDigits = resize(lonDigits)
	for i := 0; i < 9; i++ {
		bucketDigits += string(latDigits[i])
		bucketDigits += string(lonDigits[i])
	}
	return bucketDigits
}

func resize(component string) string {
	for len(component) > 9 {
		component = component[0 : len(component)-1]
	}
	for len(component) < 9 {
		component = "0" + component
	}
	return component
}
<|endoftext|>


<commit_before>package clients

import (
	"github.com/uber/zanzibar/examples/example-gateway/clients/bar"
	"github.com/uber/zanzibar/examples/example-gateway/clients/contacts"
	"github.com/uber/zanzibar/examples/example-gateway/clients/googlenow"
	"github.com/uber/zanzibar/runtime"
)

// Clients datastructure that holds all the generated clients
// This should only hold clients generate from specs
type Clients struct {
	Contacts  *contactsClient.ContactsClient
	GoogleNow *googlenowClient.GoogleNowClient
	Bar       *barClient.BarClient
}

// CreateClients will make all clients
func CreateClients(config *zanzibar.StaticConfig) *Clients {
	return &Clients{
		Contacts:  contactsClient.NewClient(config),
		GoogleNow: googlenowClient.NewClient(config),
		Bar:       barClient.NewClient(config),
	}
}
<commit_msg>example-gateway: generate client init code<commit_after>// Code generated by zanzibar
// @generated

package clients

import (
	"github.com/uber/zanzibar/examples/example-gateway/clients/bar"
	"github.com/uber/zanzibar/examples/example-gateway/clients/contacts"
	"github.com/uber/zanzibar/examples/example-gateway/clients/googlenow"

	"github.com/uber/zanzibar/runtime"
)

// Clients datastructure that holds all the generated clients
// This should only hold clients generate from specs
type Clients struct {
	GoogleNow *googlenowClient.GoogleNowClient
	Bar       *barClient.BarClient
	Contacts  *contactsClient.ContactsClient
}

// CreateClients will make all clients
func CreateClients(config *zanzibar.StaticConfig) *Clients {
	return &Clients{
		GoogleNow: googlenowClient.NewClient(config),
		Bar:       barClient.NewClient(config),
		Contacts:  contactsClient.NewClient(config),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/fsouza/go-dockerclient"

	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type Image struct {
	Id          string
	ParentId    string   `json:",omitempty"`
	RepoTags    []string `json:",omitempty"`
	VirtualSize int64
	Size        int64
	Created     int64
}

type ImagesCommand struct {
	Dot         bool `short:"d" long:"dot" description:"Show image information as Graphviz dot."`
	Tree        bool `short:"t" long:"tree" description:"Show image information as tree."`
	Short       bool `short:"s" long:"short" description:"Show short summary of images (repo name and list of tags)."`
	NoTruncate  bool `short:"n" long:"no-trunc" description:"Don't truncate the image IDs."`
	OnlyLabeled bool `short:"l" long:"only-labeled" description:"Print only labeled images/containers."`
}

var imagesCommand ImagesCommand

func (x *ImagesCommand) Execute(args []string) error {
	var images *[]Image

	stat, err := os.Stdin.Stat()
	if err != nil {
		return fmt.Errorf("error reading stdin stat", err)
	}

	if (stat.Mode() & os.ModeCharDevice) == 0 {
		// read in stdin
		stdin, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			return fmt.Errorf("error reading all input", err)
		}

		images, err = parseImagesJSON(stdin)
		if err != nil {
			return err
		}

	} else {

		client, err := connect()
		if err != nil {
			return err
		}

		clientImages, err := client.ListImages(docker.ListImagesOptions{All: true})
		if err != nil {
			if in_docker := os.Getenv("IN_DOCKER"); len(in_docker) > 0 {
				return fmt.Errorf("Unable to access Docker socket, please run like this:\n  docker run --rm -v /var/run/docker.sock:/var/run/docker.sock nate/dockviz images <args>\nFor more help, run 'dockviz help'")
			} else {
				return fmt.Errorf("Unable to connect: %s\nFor help, run 'dockviz help'", err)
			}
		}

		var ims []Image
		for _, image := range clientImages {
			// fmt.Println(image)
			ims = append(ims, Image{
				image.ID,
				image.ParentID,
				image.RepoTags,
				image.VirtualSize,
				image.Size,
				image.Created,
			})
		}

		images = &ims
	}

	 if imagesCommand.Tree || imagesCommand.Dot {
		var startImage *Image
		if len(args) > 0 {

			// attempt to find the start image, which can be specified as an
			// image ID or a repository name
			startImageArg := args[0]
			startImageRepo := args[0]

			// if tag is not defined, find by :latest tag
			if strings.Index(startImageRepo, ":") == -1 {
				startImageRepo = fmt.Sprintf("%s:latest", startImageRepo)
			}

		IMAGES:
			for _, image := range *images {
				// find by image id
				if strings.Index(image.Id, startImageArg) == 0 {
					startImage = &image
					break IMAGES
				}

				// find by image name (name and tag)
				for _, repotag := range image.RepoTags {
					if repotag == startImageRepo {
						startImage = &image
						break IMAGES
					}
				}
			}

			if startImage == nil {
				return fmt.Errorf("Unable to find image %s = %s.", startImageArg, startImageRepo)
			}
		}

		// select the start image of the tree
		var roots []Image
		if startImage == nil {
			roots = collectRoots(images)
		} else {
			startImage.ParentId = ""
			roots = []Image{*startImage}
		}

		// build helper map (image -> children)
		var imagesByParent = make(map[string][]Image)
		imagesByParent = collectChildren(images);

		// image ids truncate
		// initialize image informations

		// filter images
		if imagesCommand.OnlyLabeled{
			*images, imagesByParent = filterImages(images, &imagesByParent)
		}

		var buffer bytes.Buffer

		if imagesCommand.Tree {
			jsonToText(&buffer, imagesCommand.NoTruncate, roots, imagesByParent, "")
		}
		if imagesCommand.Dot {
			buffer.WriteString("digraph docker {\n")
			imagesToDot(&buffer, roots, imagesByParent)
			buffer.WriteString(" base [style=invisible]\n}\n")
			buffer.String()
		}

		fmt.Print(buffer.String())
	} else if imagesCommand.Short {
		fmt.Printf(jsonToShort(images))
	} else {
		return fmt.Errorf("Please specify either --dot, --tree, or --short")
	}

	return nil
}

func collectChildren(images *[]Image) map[string][]Image {
	var imagesByParent = make(map[string][]Image)
	for _, image := range *images {
		if children, exists := imagesByParent[image.ParentId]; exists {
			imagesByParent[image.ParentId] = append(children, image)
		} else {
			imagesByParent[image.ParentId] = []Image{image}
		}
	}

	return imagesByParent
}

func collectRoots(images *[]Image) []Image {
	var roots []Image
	for _, image := range *images {
		if image.ParentId == "" {
			roots = append(roots, image)
		}
	}

	return roots
}

func filterImages (images *[]Image, byParent *map[string][]Image) (filteredImages []Image, filteredChildren map[string][]Image) {
	for i := 0; i<len(*images); i++ {
		// image is visible
		//   1. it has a label
		//   2. it is root
		//   3. it is a node
		var visible bool = (*images)[i].RepoTags[0] != "<none>:<none>" || (*images)[i].ParentId == "" || len((*byParent)[(*images)[i].Id]) > 1
		if visible {
			filteredImages = append(filteredImages, (*images)[i])
		} else {
			// change childs parent id
			// if items are filtered with only one child
			for j := 0; j<len(filteredImages); j++ {
				if filteredImages[j].ParentId == (*images)[i].Id {
					filteredImages[j].ParentId = (*images)[i].ParentId
				}
			}
			for j := 0; j<len(*images); j++ {
				if (*images)[j].ParentId == (*images)[i].Id {
					(*images)[j].ParentId = (*images)[i].ParentId
				}
			}
		}
	}

	filteredChildren = collectChildren(&filteredImages)

	return filteredImages, filteredChildren
}

func jsonToText(buffer *bytes.Buffer, noTrunc bool, images []Image, byParent map[string][]Image, prefix string) {
	var length = len(images)
	if length > 1 {
		for index, image := range images {
			var nextPrefix string = ""
			if index+1 == length {
				PrintTreeNode(buffer, noTrunc, image, prefix+"└─")
				nextPrefix = "  "
			} else {
				PrintTreeNode(buffer, noTrunc, image, prefix+"├─")
				nextPrefix = "│ "
			}
			if subimages, exists := byParent[image.Id]; exists {
				jsonToText(buffer, noTrunc, subimages, byParent, prefix+nextPrefix)
			}
		}
	} else {
		for _, image := range images {
			PrintTreeNode(buffer, noTrunc, image, prefix+"└─")
			if subimages, exists := byParent[image.Id]; exists {
				jsonToText(buffer, noTrunc, subimages, byParent, prefix+"  ")
			}
		}
	}
}

func PrintTreeNode(buffer *bytes.Buffer, noTrunc bool, image Image, prefix string) {
	var imageID string
	if noTrunc {
		imageID = image.Id
	} else {
		imageID = truncate(image.Id)
	}

	buffer.WriteString(fmt.Sprintf("%s%s Virtual Size: %s", prefix, imageID, humanSize(image.VirtualSize)))
	if image.RepoTags[0] != "<none>:<none>" {
		buffer.WriteString(fmt.Sprintf(" Tags: %s\n", strings.Join(image.RepoTags, ", ")))
	} else {
		buffer.WriteString(fmt.Sprintf("\n"))
	}
}

func humanSize(raw int64) string {
	sizes := []string{"B", "KB", "MB", "GB", "TB"}

	rawFloat := float64(raw)
	ind := 0

	for {
		if rawFloat < 1000 {
			break
		} else {
			rawFloat = rawFloat / 1000
			ind = ind + 1
		}
	}

	return fmt.Sprintf("%.01f %s", rawFloat, sizes[ind])
}

func truncate(id string) string {
	return id[0:12]
}

func parseImagesJSON(rawJSON []byte) (*[]Image, error) {

	var images []Image
	err := json.Unmarshal(rawJSON, &images)

	if err != nil {
		return nil, fmt.Errorf("Error reading JSON: ", err)
	}

	return &images, nil
}

func imagesToDot(buffer *bytes.Buffer, images []Image, byParent map[string][]Image) {
	for _, image := range images {
		if image.ParentId == "" {
			buffer.WriteString(fmt.Sprintf(" base -> \"%s\" [style=invis]\n", truncate(image.Id)))
		} else {
			buffer.WriteString(fmt.Sprintf(" \"%s\" -> \"%s\"\n", truncate(image.ParentId), truncate(image.Id)))
		}
		if image.RepoTags[0] != "<none>:<none>" {
			buffer.WriteString(fmt.Sprintf(" \"%s\" [label=\"%s\\n%s\",shape=box,fillcolor=\"paleturquoise\",style=\"filled,rounded\"];\n", truncate(image.Id), truncate(image.Id), strings.Join(image.RepoTags, "\\n")))
		}
		if subimages, exists := byParent[image.Id]; exists {
			imagesToDot(buffer, subimages, byParent)
		}
	}
}

func jsonToShort(images *[]Image) string {
	var buffer bytes.Buffer

	var byRepo = make(map[string][]string)

	for _, image := range *images {
		for _, repotag := range image.RepoTags {
			if repotag != "<none>:<none>" {

				// parse the repo name and tag name out
				// tag is after the last colon
				lastColonIndex := strings.LastIndex(repotag, ":")
				tagname := repotag[lastColonIndex+1:]
				reponame := repotag[0:lastColonIndex]

				if tags, exists := byRepo[reponame]; exists {
					byRepo[reponame] = append(tags, tagname)
				} else {
					byRepo[reponame] = []string{tagname}
				}
			}
		}
	}

	for repo, tags := range byRepo {
		buffer.WriteString(fmt.Sprintf("%s: %s\n", repo, strings.Join(tags, ", ")))
	}

	return buffer.String()
}

func init() {
	parser.AddCommand("images",
		"Visualize docker images.",
		"",
		&imagesCommand)
}
<commit_msg>Fix text<commit_after>package main

import (
	"github.com/fsouza/go-dockerclient"

	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type Image struct {
	Id          string
	ParentId    string   `json:",omitempty"`
	RepoTags    []string `json:",omitempty"`
	VirtualSize int64
	Size        int64
	Created     int64
}

type ImagesCommand struct {
	Dot          bool `short:"d" long:"dot" description:"Show image information as Graphviz dot. You can add a start image id or name -d/--dot [id/name]"`
	Tree         bool `short:"t" long:"tree" description:"Show image information as tree. You can add a start image id or name -t/--tree [id/name]"`
	Short        bool `short:"s" long:"short" description:"Show short summary of images (repo name and list of tags)."`
	NoTruncate   bool `short:"n" long:"no-trunc" description:"Don't truncate the image IDs."`
	OnlyLabelled bool `short:"l" long:"only-labelled" description:"Print only labelled images/containers."`
}

var imagesCommand ImagesCommand

func (x *ImagesCommand) Execute(args []string) error {
	var images *[]Image

	stat, err := os.Stdin.Stat()
	if err != nil {
		return fmt.Errorf("error reading stdin stat", err)
	}

	if (stat.Mode() & os.ModeCharDevice) == 0 {
		// read in stdin
		stdin, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			return fmt.Errorf("error reading all input", err)
		}

		images, err = parseImagesJSON(stdin)
		if err != nil {
			return err
		}

	} else {

		client, err := connect()
		if err != nil {
			return err
		}

		clientImages, err := client.ListImages(docker.ListImagesOptions{All: true})
		if err != nil {
			if in_docker := os.Getenv("IN_DOCKER"); len(in_docker) > 0 {
				return fmt.Errorf("Unable to access Docker socket, please run like this:\n  docker run --rm -v /var/run/docker.sock:/var/run/docker.sock nate/dockviz images <args>\nFor more help, run 'dockviz help'")
			} else {
				return fmt.Errorf("Unable to connect: %s\nFor help, run 'dockviz help'", err)
			}
		}

		var ims []Image
		for _, image := range clientImages {
			// fmt.Println(image)
			ims = append(ims, Image{
				image.ID,
				image.ParentID,
				image.RepoTags,
				image.VirtualSize,
				image.Size,
				image.Created,
			})
		}

		images = &ims
	}

	 if imagesCommand.Tree || imagesCommand.Dot {
		var startImage *Image
		if len(args) > 0 {

			// attempt to find the start image, which can be specified as an
			// image ID or a repository name
			startImageArg := args[0]
			startImageRepo := args[0]

			// if tag is not defined, find by :latest tag
			if strings.Index(startImageRepo, ":") == -1 {
				startImageRepo = fmt.Sprintf("%s:latest", startImageRepo)
			}

		IMAGES:
			for _, image := range *images {
				// find by image id
				if strings.Index(image.Id, startImageArg) == 0 {
					startImage = &image
					break IMAGES
				}

				// find by image name (name and tag)
				for _, repotag := range image.RepoTags {
					if repotag == startImageRepo {
						startImage = &image
						break IMAGES
					}
				}
			}

			if startImage == nil {
				return fmt.Errorf("Unable to find image %s = %s.", startImageArg, startImageRepo)
			}
		}

		// select the start image of the tree
		var roots []Image
		if startImage == nil {
			roots = collectRoots(images)
		} else {
			startImage.ParentId = ""
			roots = []Image{*startImage}
		}

		// build helper map (image -> children)
		var imagesByParent = make(map[string][]Image)
		imagesByParent = collectChildren(images);

		// image ids truncate
		// initialize image informations

		// filter images
		if imagesCommand.OnlyLabelled{
			*images, imagesByParent = filterImages(images, &imagesByParent)
		}

		var buffer bytes.Buffer

		if imagesCommand.Tree {
			jsonToText(&buffer, imagesCommand.NoTruncate, roots, imagesByParent, "")
		}
		if imagesCommand.Dot {
			buffer.WriteString("digraph docker {\n")
			imagesToDot(&buffer, roots, imagesByParent)
			buffer.WriteString(" base [style=invisible]\n}\n")
			buffer.String()
		}

		fmt.Print(buffer.String())
	} else if imagesCommand.Short {
		fmt.Printf(jsonToShort(images))
	} else {
		return fmt.Errorf("Please specify either --dot, --tree, or --short")
	}

	return nil
}

func collectChildren(images *[]Image) map[string][]Image {
	var imagesByParent = make(map[string][]Image)
	for _, image := range *images {
		if children, exists := imagesByParent[image.ParentId]; exists {
			imagesByParent[image.ParentId] = append(children, image)
		} else {
			imagesByParent[image.ParentId] = []Image{image}
		}
	}

	return imagesByParent
}

func collectRoots(images *[]Image) []Image {
	var roots []Image
	for _, image := range *images {
		if image.ParentId == "" {
			roots = append(roots, image)
		}
	}

	return roots
}

func filterImages (images *[]Image, byParent *map[string][]Image) (filteredImages []Image, filteredChildren map[string][]Image) {
	for i := 0; i<len(*images); i++ {
		// image is visible
		//   1. it has a label
		//   2. it is root
		//   3. it is a node
		var visible bool = (*images)[i].RepoTags[0] != "<none>:<none>" || (*images)[i].ParentId == "" || len((*byParent)[(*images)[i].Id]) > 1
		if visible {
			filteredImages = append(filteredImages, (*images)[i])
		} else {
			// change childs parent id
			// if items are filtered with only one child
			for j := 0; j<len(filteredImages); j++ {
				if filteredImages[j].ParentId == (*images)[i].Id {
					filteredImages[j].ParentId = (*images)[i].ParentId
				}
			}
			for j := 0; j<len(*images); j++ {
				if (*images)[j].ParentId == (*images)[i].Id {
					(*images)[j].ParentId = (*images)[i].ParentId
				}
			}
		}
	}

	filteredChildren = collectChildren(&filteredImages)

	return filteredImages, filteredChildren
}

func jsonToText(buffer *bytes.Buffer, noTrunc bool, images []Image, byParent map[string][]Image, prefix string) {
	var length = len(images)
	if length > 1 {
		for index, image := range images {
			var nextPrefix string = ""
			if index+1 == length {
				PrintTreeNode(buffer, noTrunc, image, prefix+"└─")
				nextPrefix = "  "
			} else {
				PrintTreeNode(buffer, noTrunc, image, prefix+"├─")
				nextPrefix = "│ "
			}
			if subimages, exists := byParent[image.Id]; exists {
				jsonToText(buffer, noTrunc, subimages, byParent, prefix+nextPrefix)
			}
		}
	} else {
		for _, image := range images {
			PrintTreeNode(buffer, noTrunc, image, prefix+"└─")
			if subimages, exists := byParent[image.Id]; exists {
				jsonToText(buffer, noTrunc, subimages, byParent, prefix+"  ")
			}
		}
	}
}

func PrintTreeNode(buffer *bytes.Buffer, noTrunc bool, image Image, prefix string) {
	var imageID string
	if noTrunc {
		imageID = image.Id
	} else {
		imageID = truncate(image.Id)
	}

	buffer.WriteString(fmt.Sprintf("%s%s Virtual Size: %s", prefix, imageID, humanSize(image.VirtualSize)))
	if image.RepoTags[0] != "<none>:<none>" {
		buffer.WriteString(fmt.Sprintf(" Tags: %s\n", strings.Join(image.RepoTags, ", ")))
	} else {
		buffer.WriteString(fmt.Sprintf("\n"))
	}
}

func humanSize(raw int64) string {
	sizes := []string{"B", "KB", "MB", "GB", "TB"}

	rawFloat := float64(raw)
	ind := 0

	for {
		if rawFloat < 1000 {
			break
		} else {
			rawFloat = rawFloat / 1000
			ind = ind + 1
		}
	}

	return fmt.Sprintf("%.01f %s", rawFloat, sizes[ind])
}

func truncate(id string) string {
	return id[0:12]
}

func parseImagesJSON(rawJSON []byte) (*[]Image, error) {

	var images []Image
	err := json.Unmarshal(rawJSON, &images)

	if err != nil {
		return nil, fmt.Errorf("Error reading JSON: ", err)
	}

	return &images, nil
}

func imagesToDot(buffer *bytes.Buffer, images []Image, byParent map[string][]Image) {
	for _, image := range images {
		if image.ParentId == "" {
			buffer.WriteString(fmt.Sprintf(" base -> \"%s\" [style=invis]\n", truncate(image.Id)))
		} else {
			buffer.WriteString(fmt.Sprintf(" \"%s\" -> \"%s\"\n", truncate(image.ParentId), truncate(image.Id)))
		}
		if image.RepoTags[0] != "<none>:<none>" {
			buffer.WriteString(fmt.Sprintf(" \"%s\" [label=\"%s\\n%s\",shape=box,fillcolor=\"paleturquoise\",style=\"filled,rounded\"];\n", truncate(image.Id), truncate(image.Id), strings.Join(image.RepoTags, "\\n")))
		}
		if subimages, exists := byParent[image.Id]; exists {
			imagesToDot(buffer, subimages, byParent)
		}
	}
}

func jsonToShort(images *[]Image) string {
	var buffer bytes.Buffer

	var byRepo = make(map[string][]string)

	for _, image := range *images {
		for _, repotag := range image.RepoTags {
			if repotag != "<none>:<none>" {

				// parse the repo name and tag name out
				// tag is after the last colon
				lastColonIndex := strings.LastIndex(repotag, ":")
				tagname := repotag[lastColonIndex+1:]
				reponame := repotag[0:lastColonIndex]

				if tags, exists := byRepo[reponame]; exists {
					byRepo[reponame] = append(tags, tagname)
				} else {
					byRepo[reponame] = []string{tagname}
				}
			}
		}
	}

	for repo, tags := range byRepo {
		buffer.WriteString(fmt.Sprintf("%s: %s\n", repo, strings.Join(tags, ", ")))
	}

	return buffer.String()
}

func init() {
	parser.AddCommand("images",
		"Visualize docker images.",
		"",
		&imagesCommand)
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 cache

import (
	"context"
	"fmt"

	"github.com/google/exposure-notifications-verification-server/pkg/redis"
	"github.com/sethvargo/go-envconfig"
)

// CacherType represents a type of cacher.
type CacherType string

const (
	TypeNoop     CacherType = "NOOP"
	TypeInMemory CacherType = "IN_MEMORY"
	TypeRedis    CacherType = "REDIS"
)

// Config represents configuration for a cacher.
type Config struct {
	Type CacherType `env:"TYPE, default=IN_MEMORY"`

	// HMACKey is the hash key to use when for keys in the cacher.
	HMACKey envconfig.Base64Bytes `env:"CACHE_HMAC_KEY, required"`

	// Redis configuration
	Redis redis.Config `env:", prefix=CACHE_"`
}

func CacherFor(ctx context.Context, c *Config, keyFunc KeyFunc) (Cacher, error) {
	switch typ := c.Type; typ {
	case TypeNoop:
		return NewNoop()
	case TypeInMemory:
		return NewInMemory(&InMemoryConfig{
			KeyFunc: keyFunc,
		})
	case TypeRedis:
		return NewRedis(&RedisConfig{
			KeyFunc:     keyFunc,
			Address:     fmt.Sprintf("%s:%s", c.Redis.Host, c.Redis.Port),
			Username:    c.Redis.Username,
			Password:    c.Redis.Password,
			IdleTimeout: c.Redis.IdleTimeout,
			MaxIdle:     c.Redis.MaxIdle,
			MaxActive:   c.Redis.MaxActive,
		})
	default:
		return nil, fmt.Errorf("unknown cacher type: %v", typ)
	}
}
<commit_msg>Fix cache type envvar as CACHE_TYPE (#635)<commit_after>// Copyright 2020 Google LLC
//
// 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 cache

import (
	"context"
	"fmt"

	"github.com/google/exposure-notifications-verification-server/pkg/redis"
	"github.com/sethvargo/go-envconfig"
)

// CacherType represents a type of cacher.
type CacherType string

const (
	TypeNoop     CacherType = "NOOP"
	TypeInMemory CacherType = "IN_MEMORY"
	TypeRedis    CacherType = "REDIS"
)

// Config represents configuration for a cacher.
type Config struct {
	Type CacherType `env:"CACHE_TYPE, default=IN_MEMORY"`

	// HMACKey is the hash key to use when for keys in the cacher.
	HMACKey envconfig.Base64Bytes `env:"CACHE_HMAC_KEY, required"`

	// Redis configuration
	Redis redis.Config `env:", prefix=CACHE_"`
}

func CacherFor(ctx context.Context, c *Config, keyFunc KeyFunc) (Cacher, error) {
	switch typ := c.Type; typ {
	case TypeNoop:
		return NewNoop()
	case TypeInMemory:
		return NewInMemory(&InMemoryConfig{
			KeyFunc: keyFunc,
		})
	case TypeRedis:
		return NewRedis(&RedisConfig{
			KeyFunc:     keyFunc,
			Address:     fmt.Sprintf("%s:%s", c.Redis.Host, c.Redis.Port),
			Username:    c.Redis.Username,
			Password:    c.Redis.Password,
			IdleTimeout: c.Redis.IdleTimeout,
			MaxIdle:     c.Redis.MaxIdle,
			MaxActive:   c.Redis.MaxActive,
		})
	default:
		return nil, fmt.Errorf("unknown cacher type: %v", typ)
	}
}
<|endoftext|>


<commit_before>package model

import (
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/stretchr/testify/assert"
)

func TestSetInstance(t *testing.T) {
	i := NewInstance(&ec2.Instance{
		InstanceId: aws.String("i-1000000"),
	})
	eni := NewENI(&ec2.NetworkInterface{})

	eni.SetInstance(i)

	assert.Equal(t, eni.instance, i)
}

func TestInterfaceID(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
	})

	assert.Equal(t, eni.InterfaceID(), "eni-2222222")
}

func TestPrivateDnsName(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		PrivateDnsName: aws.String("ip-10-0-0-100.ap-northeast-1.compute.internal"),
	})

	assert.Equal(t, eni.PrivateDnsName(), "ip-10-0-0-100.ap-northeast-1.compute.internal")
}

func TestPrivateIpAddress(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		PrivateIpAddress: aws.String("10.0.0.100"),
	})

	assert.Equal(t, eni.PrivateIpAddress(), "10.0.0.100")
}

func TestStatus(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Status: aws.String("in-use"),
	})

	assert.Equal(t, eni.Status(), "in-use")
}

func TestAttachmentID(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			AttachmentId: aws.String("eni-attach-11111111"),
		},
	})

	assert.Equal(t, eni.AttachmentID(), "eni-attach-11111111")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachmentID(), "")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			AttachmentId: nil,
		},
	})

	assert.Equal(t, eni.AttachmentID(), "")
}

func TestAttachedDeviceIndex(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			DeviceIndex: aws.Int64(0),
		},
	})

	assert.Equal(t, eni.AttachedDeviceIndex(), int64(0))

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachedDeviceIndex(), int64(-1))

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			DeviceIndex: nil,
		},
	})

	assert.Equal(t, eni.AttachedDeviceIndex(), int64(-1))
}

func TestAttachedStatus(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			Status: aws.String("in-use"),
		},
	})

	assert.Equal(t, eni.AttachedStatus(), "in-use")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachedStatus(), "")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			Status: nil,
		},
	})

	assert.Equal(t, eni.AttachedStatus(), "")
}

func TestAttachedInstanceID(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			InstanceId: aws.String("i-1111111"),
		},
	})

	assert.Equal(t, eni.AttachedInstanceID(), "i-1111111")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachedInstanceID(), "")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			InstanceId: nil,
		},
	})

	assert.Equal(t, eni.AttachedInstanceID(), "")
}

func TestAttachedInstance(t *testing.T) {
	i := NewInstance(&ec2.Instance{
		InstanceId: aws.String("i-1000000"),
	})
	eni := NewENI(&ec2.NetworkInterface{})

	eni.SetInstance(i)

	assert.Equal(t, *eni.AttachedInstance().InstanceId, "i-1000000")
}

func TestAvailabilityZone(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		AvailabilityZone: aws.String("ap-northeast-1a"),
	})

	assert.Equal(t, eni.AvailabilityZone(), "ap-northeast-1a")
}
<commit_msg>Add TestInterfaceName<commit_after>package model

import (
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/stretchr/testify/assert"
)

func TestSetInstance(t *testing.T) {
	i := NewInstance(&ec2.Instance{
		InstanceId: aws.String("i-1000000"),
	})
	eni := NewENI(&ec2.NetworkInterface{})

	eni.SetInstance(i)

	assert.Equal(t, eni.instance, i)
}

func TestInterfaceID(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
	})

	assert.Equal(t, eni.InterfaceID(), "eni-2222222")
}

func TestPrivateDnsName(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		PrivateDnsName: aws.String("ip-10-0-0-100.ap-northeast-1.compute.internal"),
	})

	assert.Equal(t, eni.PrivateDnsName(), "ip-10-0-0-100.ap-northeast-1.compute.internal")
}

func TestPrivateIpAddress(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		PrivateIpAddress: aws.String("10.0.0.100"),
	})

	assert.Equal(t, eni.PrivateIpAddress(), "10.0.0.100")
}

func TestStatus(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Status: aws.String("in-use"),
	})

	assert.Equal(t, eni.Status(), "in-use")
}

func TestAttachmentID(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			AttachmentId: aws.String("eni-attach-11111111"),
		},
	})

	assert.Equal(t, eni.AttachmentID(), "eni-attach-11111111")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachmentID(), "")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			AttachmentId: nil,
		},
	})

	assert.Equal(t, eni.AttachmentID(), "")
}

func TestAttachedDeviceIndex(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			DeviceIndex: aws.Int64(0),
		},
	})

	assert.Equal(t, eni.AttachedDeviceIndex(), int64(0))

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachedDeviceIndex(), int64(-1))

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			DeviceIndex: nil,
		},
	})

	assert.Equal(t, eni.AttachedDeviceIndex(), int64(-1))
}

func TestAttachedStatus(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			Status: aws.String("in-use"),
		},
	})

	assert.Equal(t, eni.AttachedStatus(), "in-use")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachedStatus(), "")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			Status: nil,
		},
	})

	assert.Equal(t, eni.AttachedStatus(), "")
}

func TestAttachedInstanceID(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			InstanceId: aws.String("i-1111111"),
		},
	})

	assert.Equal(t, eni.AttachedInstanceID(), "i-1111111")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: nil,
	})

	assert.Equal(t, eni.AttachedInstanceID(), "")

	eni = NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		Attachment: &ec2.NetworkInterfaceAttachment{
			InstanceId: nil,
		},
	})

	assert.Equal(t, eni.AttachedInstanceID(), "")
}

func TestAttachedInstance(t *testing.T) {
	i := NewInstance(&ec2.Instance{
		InstanceId: aws.String("i-1000000"),
	})
	eni := NewENI(&ec2.NetworkInterface{})

	eni.SetInstance(i)

	assert.Equal(t, *eni.AttachedInstance().InstanceId, "i-1000000")
}

func TestAvailabilityZone(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		NetworkInterfaceId: aws.String("eni-2222222"),
		AvailabilityZone: aws.String("ap-northeast-1a"),
	})

	assert.Equal(t, eni.AvailabilityZone(), "ap-northeast-1a")
}

func TestInterfaceName(t *testing.T) {
	eni := NewENI(&ec2.NetworkInterface{
		TagSet: []*ec2.Tag{&ec2.Tag{
			Key: aws.String("Name"),
			Value: aws.String("eni001"),
		}},
	})

	assert.Equal(t, eni.Name(), "eni001")
}
<|endoftext|>


<commit_before>package redis

import (
	"bytes"
	"errors"
	"fmt"
	. "testing"

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

func pretendRead(s string) *Resp {
	buf := bytes.NewBufferString(s)
	return NewRespReader(buf).Read()
}

func TestRead(t *T) {

	r := pretendRead("")
	assert.NotNil(t, r.Err)

	// Simple string
	r = pretendRead("+ohey\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte("ohey"), r.val)
	s, err := r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "ohey", s)

	// Empty simple string
	r = pretendRead("+\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte(""), r.val)
	s, err = r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "", s)

	// Error
	r = pretendRead("-ohey\r\n")
	assert.Equal(t, AppErr, r.typ)
	assert.Exactly(t, errors.New("ohey"), r.val)
	assert.Equal(t, "ohey", r.Err.Error())

	// Empty error
	r = pretendRead("-\r\n")
	assert.Equal(t, AppErr, r.typ)
	assert.Exactly(t, errors.New(""), r.val)
	assert.Equal(t, "", r.Err.Error())

	// Int
	r = pretendRead(":1024\r\n")
	assert.Equal(t, Int, r.typ)
	assert.Exactly(t, int64(1024), r.val)
	i, err := r.Int()
	assert.Nil(t, err)
	assert.Equal(t, 1024, i)

	// Int (from string)
	r = pretendRead("+50\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte("50"), r.val)
	i, err = r.Int()
	assert.Nil(t, err)
	assert.Equal(t, 50, i)

	// Int (from string, can't parse)
	r = pretendRead("+ImADuck\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte("ImADuck"), r.val)
	i, err = r.Int()
	assert.NotNil(t, err)

	// Bulk string
	r = pretendRead("$3\r\nfoo\r\n")
	assert.Equal(t, BulkStr, r.typ)
	assert.Exactly(t, []byte("foo"), r.val)
	s, err = r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "foo", s)

	// Empty bulk string
	r = pretendRead("$0\r\n\r\n")
	assert.Equal(t, BulkStr, r.typ)
	assert.Exactly(t, []byte(""), r.val)
	s, err = r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "", s)

	// Nil bulk string
	r = pretendRead("$-1\r\n")
	assert.Equal(t, Nil, r.typ)

	// Array
	r = pretendRead("*2\r\n+foo\r\n+bar\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 2, len(r.val.([]Resp)))
	assert.Equal(t, SimpleStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, SimpleStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)
	l, err := r.List()
	assert.Nil(t, err)
	assert.Equal(t, []string{"foo", "bar"}, l)
	b, err := r.ListBytes()
	assert.Nil(t, err)
	assert.Equal(t, [][]byte{[]byte("foo"), []byte("bar")}, b)
	m, err := r.Map()
	assert.Nil(t, err)
	assert.Equal(t, map[string]string{"foo": "bar"}, m)

	// Empty Array
	r = pretendRead("*0\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 0, len(r.val.([]Resp)))

	// Nil Array
	r = pretendRead("*-1\r\n")
	assert.Equal(t, Nil, r.typ)

	// Embedded Array
	r = pretendRead("*3\r\n+foo\r\n+bar\r\n*2\r\n+foo\r\n+bar\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 3, len(r.val.([]Resp)))
	assert.Equal(t, SimpleStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, SimpleStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)
	r = &r.val.([]Resp)[2]
	assert.Equal(t, 2, len(r.val.([]Resp)))
	assert.Equal(t, SimpleStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, SimpleStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)

	// Test that two bulks in a row read correctly
	r = pretendRead("*2\r\n$3\r\nfoo\r\n$3\r\nbar\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 2, len(r.val.([]Resp)))
	assert.Equal(t, BulkStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, BulkStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)
}

type arbitraryTest struct {
	val    interface{}
	expect []byte
}

var nilResp = pretendRead("$-1\r\n")

var arbitraryTests = []arbitraryTest{
	{[]byte("OHAI"), []byte("$4\r\nOHAI\r\n")},
	{"OHAI", []byte("$4\r\nOHAI\r\n")},
	{true, []byte("$1\r\n1\r\n")},
	{false, []byte("$1\r\n0\r\n")},
	{nil, []byte("$-1\r\n")},
	{80, []byte(":80\r\n")},
	{int64(-80), []byte(":-80\r\n")},
	{uint64(80), []byte(":80\r\n")},
	{float32(0.1234), []byte("$6\r\n0.1234\r\n")},
	{float64(0.1234), []byte("$6\r\n0.1234\r\n")},
	{errors.New("hi"), []byte("-hi\r\n")},

	{nilResp, []byte("$-1\r\n")},

	{[]int{1, 2, 3}, []byte("*3\r\n:1\r\n:2\r\n:3\r\n")},
	{map[int]int{1: 2}, []byte("*2\r\n:1\r\n:2\r\n")},

	{NewRespSimple("OK"), []byte("+OK\r\n")},
}

var arbitraryAsStringTests = []arbitraryTest{
	{[]byte("OHAI"), []byte("$4\r\nOHAI\r\n")},
	{"OHAI", []byte("$4\r\nOHAI\r\n")},
	{true, []byte("$1\r\n1\r\n")},
	{false, []byte("$1\r\n0\r\n")},
	{nil, []byte("$0\r\n\r\n")},
	{80, []byte("$2\r\n80\r\n")},
	{int64(-80), []byte("$3\r\n-80\r\n")},
	{uint64(80), []byte("$2\r\n80\r\n")},
	{float32(0.1234), []byte("$6\r\n0.1234\r\n")},
	{float64(0.1234), []byte("$6\r\n0.1234\r\n")},
	{errors.New("hi"), []byte("$2\r\nhi\r\n")},

	{nilResp, []byte("$-1\r\n")},

	{[]int{1, 2, 3}, []byte("*3\r\n$1\r\n1\r\n$1\r\n2\r\n$1\r\n3\r\n")},
	{map[int]int{1: 2}, []byte("*2\r\n$1\r\n1\r\n$1\r\n2\r\n")},

	{NewRespSimple("OK"), []byte("+OK\r\n")},
}

var arbitraryAsFlattenedStringsTests = []arbitraryTest{
	{
		[]interface{}{"wat", map[string]interface{}{
			"foo": 1,
		}},
		[]byte("*3\r\n$3\r\nwat\r\n$3\r\nfoo\r\n$1\r\n1\r\n"),
	},
}

func TestWriteArbitrary(t *T) {
	var err error
	buf := bytes.NewBuffer([]byte{})
	for _, test := range arbitraryTests {
		buf.Reset()
		_, err = NewResp(test.val).WriteTo(buf)
		assert.Nil(t, err)
		assert.Equal(t, test.expect, buf.Bytes())
	}
}

func TestWriteArbitraryAsFlattenedStrings(t *T) {
	var err error
	buf := bytes.NewBuffer([]byte{})
	for _, test := range arbitraryAsFlattenedStringsTests {
		buf.Reset()
		_, err = NewRespFlattenedStrings(test.val).WriteTo(buf)
		assert.Nil(t, err)
		assert.Equal(t, test.expect, buf.Bytes())
	}
}

func TestFloat64(t *T) {
	r := NewResp(4)
	_, err := r.Float64()
	assert.NotNil(t, err)

	testErr := fmt.Errorf("test")
	r = NewResp(testErr)
	_, err = r.Float64()
	assert.Equal(t, testErr, err)

	r = NewResp("test")
	_, err = r.Float64()
	assert.NotNil(t, err)

	r = NewResp("5.0")
	f, err := r.Float64()
	assert.Nil(t, err)
	assert.Equal(t, float64(5.0), f)

}
<commit_msg>remove unnecessary map from resp tests<commit_after>package redis

import (
	"bytes"
	"errors"
	"fmt"
	. "testing"

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

func pretendRead(s string) *Resp {
	buf := bytes.NewBufferString(s)
	return NewRespReader(buf).Read()
}

func TestRead(t *T) {

	r := pretendRead("")
	assert.NotNil(t, r.Err)

	// Simple string
	r = pretendRead("+ohey\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte("ohey"), r.val)
	s, err := r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "ohey", s)

	// Empty simple string
	r = pretendRead("+\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte(""), r.val)
	s, err = r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "", s)

	// Error
	r = pretendRead("-ohey\r\n")
	assert.Equal(t, AppErr, r.typ)
	assert.Exactly(t, errors.New("ohey"), r.val)
	assert.Equal(t, "ohey", r.Err.Error())

	// Empty error
	r = pretendRead("-\r\n")
	assert.Equal(t, AppErr, r.typ)
	assert.Exactly(t, errors.New(""), r.val)
	assert.Equal(t, "", r.Err.Error())

	// Int
	r = pretendRead(":1024\r\n")
	assert.Equal(t, Int, r.typ)
	assert.Exactly(t, int64(1024), r.val)
	i, err := r.Int()
	assert.Nil(t, err)
	assert.Equal(t, 1024, i)

	// Int (from string)
	r = pretendRead("+50\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte("50"), r.val)
	i, err = r.Int()
	assert.Nil(t, err)
	assert.Equal(t, 50, i)

	// Int (from string, can't parse)
	r = pretendRead("+ImADuck\r\n")
	assert.Equal(t, SimpleStr, r.typ)
	assert.Exactly(t, []byte("ImADuck"), r.val)
	i, err = r.Int()
	assert.NotNil(t, err)

	// Bulk string
	r = pretendRead("$3\r\nfoo\r\n")
	assert.Equal(t, BulkStr, r.typ)
	assert.Exactly(t, []byte("foo"), r.val)
	s, err = r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "foo", s)

	// Empty bulk string
	r = pretendRead("$0\r\n\r\n")
	assert.Equal(t, BulkStr, r.typ)
	assert.Exactly(t, []byte(""), r.val)
	s, err = r.Str()
	assert.Nil(t, err)
	assert.Equal(t, "", s)

	// Nil bulk string
	r = pretendRead("$-1\r\n")
	assert.Equal(t, Nil, r.typ)

	// Array
	r = pretendRead("*2\r\n+foo\r\n+bar\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 2, len(r.val.([]Resp)))
	assert.Equal(t, SimpleStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, SimpleStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)
	l, err := r.List()
	assert.Nil(t, err)
	assert.Equal(t, []string{"foo", "bar"}, l)
	b, err := r.ListBytes()
	assert.Nil(t, err)
	assert.Equal(t, [][]byte{[]byte("foo"), []byte("bar")}, b)
	m, err := r.Map()
	assert.Nil(t, err)
	assert.Equal(t, map[string]string{"foo": "bar"}, m)

	// Empty Array
	r = pretendRead("*0\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 0, len(r.val.([]Resp)))

	// Nil Array
	r = pretendRead("*-1\r\n")
	assert.Equal(t, Nil, r.typ)

	// Embedded Array
	r = pretendRead("*3\r\n+foo\r\n+bar\r\n*2\r\n+foo\r\n+bar\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 3, len(r.val.([]Resp)))
	assert.Equal(t, SimpleStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, SimpleStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)
	r = &r.val.([]Resp)[2]
	assert.Equal(t, 2, len(r.val.([]Resp)))
	assert.Equal(t, SimpleStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, SimpleStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)

	// Test that two bulks in a row read correctly
	r = pretendRead("*2\r\n$3\r\nfoo\r\n$3\r\nbar\r\n")
	assert.Equal(t, Array, r.typ)
	assert.Equal(t, 2, len(r.val.([]Resp)))
	assert.Equal(t, BulkStr, r.val.([]Resp)[0].typ)
	assert.Exactly(t, []byte("foo"), r.val.([]Resp)[0].val)
	assert.Equal(t, BulkStr, r.val.([]Resp)[1].typ)
	assert.Exactly(t, []byte("bar"), r.val.([]Resp)[1].val)
}

type arbitraryTest struct {
	val    interface{}
	expect []byte
}

var nilResp = pretendRead("$-1\r\n")

var arbitraryTests = []arbitraryTest{
	{[]byte("OHAI"), []byte("$4\r\nOHAI\r\n")},
	{"OHAI", []byte("$4\r\nOHAI\r\n")},
	{true, []byte("$1\r\n1\r\n")},
	{false, []byte("$1\r\n0\r\n")},
	{nil, []byte("$-1\r\n")},
	{80, []byte(":80\r\n")},
	{int64(-80), []byte(":-80\r\n")},
	{uint64(80), []byte(":80\r\n")},
	{float32(0.1234), []byte("$6\r\n0.1234\r\n")},
	{float64(0.1234), []byte("$6\r\n0.1234\r\n")},
	{errors.New("hi"), []byte("-hi\r\n")},

	{nilResp, []byte("$-1\r\n")},

	{[]int{1, 2, 3}, []byte("*3\r\n:1\r\n:2\r\n:3\r\n")},
	{map[int]int{1: 2}, []byte("*2\r\n:1\r\n:2\r\n")},

	{NewRespSimple("OK"), []byte("+OK\r\n")},
}

var arbitraryAsFlattenedStringsTests = []arbitraryTest{
	{
		[]interface{}{"wat", map[string]interface{}{
			"foo": 1,
		}},
		[]byte("*3\r\n$3\r\nwat\r\n$3\r\nfoo\r\n$1\r\n1\r\n"),
	},
}

func TestWriteArbitrary(t *T) {
	var err error
	buf := bytes.NewBuffer([]byte{})
	for _, test := range arbitraryTests {
		buf.Reset()
		_, err = NewResp(test.val).WriteTo(buf)
		assert.Nil(t, err)
		assert.Equal(t, test.expect, buf.Bytes())
	}
}

func TestWriteArbitraryAsFlattenedStrings(t *T) {
	var err error
	buf := bytes.NewBuffer([]byte{})
	for _, test := range arbitraryAsFlattenedStringsTests {
		buf.Reset()
		_, err = NewRespFlattenedStrings(test.val).WriteTo(buf)
		assert.Nil(t, err)
		assert.Equal(t, test.expect, buf.Bytes())
	}
}

func TestFloat64(t *T) {
	r := NewResp(4)
	_, err := r.Float64()
	assert.NotNil(t, err)

	testErr := fmt.Errorf("test")
	r = NewResp(testErr)
	_, err = r.Float64()
	assert.Equal(t, testErr, err)

	r = NewResp("test")
	_, err = r.Float64()
	assert.NotNil(t, err)

	r = NewResp("5.0")
	f, err := r.Float64()
	assert.Nil(t, err)
	assert.Equal(t, float64(5.0), f)

}
<|endoftext|>


<commit_before>package cuckoofilter_test

import (
	"testing"

	cuckoo2 "github.com/seiflotfy/cuckoofilter"
	cuckoo3 "github.com/tylertreat/BoomFilters"
)

func BenchmarkExt1Insert(b *testing.B) {
	j, n := 0, len(data)
	f := cuckoo2.NewCuckooFilter(uint(n))
	for i := 0; i < b.N; i++ {
		f.Insert(data[j])
		if j++; j == n {
			j = 0
		}
	}
}

func BenchmarkExt1Has(b *testing.B) {
	f := *fullFilter2
	j, n := 0, int(f.GetCount())
	for i := 0; i < b.N; i++ {
		f.Lookup(data[j])
		if j++; j == n {
			j = 0
		}
	}
}

func BenchmarkExt2Insert(b *testing.B) {
	j, n := 0, len(data)
	f := cuckoo3.NewCuckooFilter(uint(n), 0.01)
	for i := 0; i < b.N; i++ {
		f.Add(data[j])
		if j++; j == n {
			j = 0
		}
	}
}

func BenchmarkExt2Has(b *testing.B) {
	f := *fullFilter3
	j, n := 0, int(f.Count())
	for i := 0; i < b.N; i++ {
		f.Test(data[j])
		if j++; j == n {
			j = 0
		}
	}
}
<commit_msg>github.com/seiflotfy/cuckoofilter.GetCount to Count<commit_after>package cuckoofilter_test

import (
	"testing"

	cuckoo2 "github.com/seiflotfy/cuckoofilter"
	cuckoo3 "github.com/tylertreat/BoomFilters"
)

func BenchmarkExt1Insert(b *testing.B) {
	j, n := 0, len(data)
	f := cuckoo2.NewCuckooFilter(uint(n))
	for i := 0; i < b.N; i++ {
		f.Insert(data[j])
		if j++; j == n {
			j = 0
		}
	}
}

func BenchmarkExt1Has(b *testing.B) {
	f := *fullFilter2
	j, n := 0, int(f.Count())
	for i := 0; i < b.N; i++ {
		f.Lookup(data[j])
		if j++; j == n {
			j = 0
		}
	}
}

func BenchmarkExt2Insert(b *testing.B) {
	j, n := 0, len(data)
	f := cuckoo3.NewCuckooFilter(uint(n), 0.01)
	for i := 0; i < b.N; i++ {
		f.Add(data[j])
		if j++; j == n {
			j = 0
		}
	}
}

func BenchmarkExt2Has(b *testing.B) {
	f := *fullFilter3
	j, n := 0, int(f.Count())
	for i := 0; i < b.N; i++ {
		f.Test(data[j])
		if j++; j == n {
			j = 0
		}
	}
}
<|endoftext|>


<commit_before>package http

import (
	"net"
	"net/http"

	"github.com/putdotio/putio-sync/sync"
)

const (
	defaultAddr = ":3000"
	debugAddr   = ":3001"
)

// Server represent the HTTP interface to the sync client.
type Server struct {
	ln      net.Listener
	Handler *Handler
	Addr    string
}

// NewServer returns a new instance of Server.
func NewServer(sync *sync.Client) *Server {
	s := &Server{
		Handler: NewHandler(sync),
	}

	if sync.Debug {
		s.Addr = debugAddr
	} else {
		s.Addr = defaultAddr
	}
	return s
}

// Open opens up the underlying socket for the HTTP server.
func (s *Server) Open() error {
	ln, err := net.Listen("tcp", s.Addr)
	if err != nil {
		return err
	}
	s.ln = ln
	return nil
}

// Close closes the underlying socket.
func (s *Server) Close() error {
	if s.ln == nil {
		return nil
	}

	return s.ln.Close()
}

// Port returns the port that the server is open on. Only valid after open.
func (s *Server) Port() int {
	return s.ln.Addr().(*net.TCPAddr).Port
}

// Serve serves the HTTP requests.
func (s *Server) Serve() error {
	// remember the last paused state and continue syncing
	if !s.Handler.sync.Config.IsPaused {
		err := s.Handler.sync.Run()
		if err != nil {
			return err
		}
	}

	return http.Serve(s.ln, s.Handler)
}
<commit_msg>sync: do not change port when debug<commit_after>package http

import (
	"net"
	"net/http"

	"github.com/putdotio/putio-sync/sync"
)

const (
	defaultAddr = ":3000"
)

// Server represent the HTTP interface to the sync client.
type Server struct {
	ln      net.Listener
	Handler *Handler
	Addr    string
}

// NewServer returns a new instance of Server.
func NewServer(sync *sync.Client) *Server {
	return &Server{
		Handler: NewHandler(sync),
		Addr:    defaultAddr,
	}
}

// Open opens up the underlying socket for the HTTP server.
func (s *Server) Open() error {
	ln, err := net.Listen("tcp", s.Addr)
	if err != nil {
		return err
	}
	s.ln = ln
	return nil
}

// Close closes the underlying socket.
func (s *Server) Close() error {
	if s.ln == nil {
		return nil
	}

	return s.ln.Close()
}

// Port returns the port that the server is open on. Only valid after open.
func (s *Server) Port() int {
	return s.ln.Addr().(*net.TCPAddr).Port
}

// Serve serves the HTTP requests.
func (s *Server) Serve() error {
	// remember the last paused state and continue syncing
	if !s.Handler.sync.Config.IsPaused {
		err := s.Handler.sync.Run()
		if err != nil {
			return err
		}
	}

	return http.Serve(s.ln, s.Handler)
}
<|endoftext|>


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

package ht

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
)

func init() {
	RegisterCheck(&Logfile{})
}

// ----------------------------------------------------------------------------
// Logfile

// Logfile provides checks on files (i.e. it ignores the response).
// During preparation the current file size is determined and the checks
// are run against the bytes written after preparation.
type Logfile struct {
	// Path is the file system path to the logfile."
	Path string

	// Condition the written stuff must fulfill.
	Condition `json:",omitempty"`

	// Disallow states what is forbidden in the written log.
	Disallow []string `json:",omitempty"`

	pos int64
}

// Execute implements Check's Execute method.
func (f *Logfile) Execute(t *Test) error {
	file, err := os.Open(f.Path)
	if err != nil {
		return err
	}
	pos, err := file.Seek(f.pos, os.SEEK_SET)
	if err != nil {
		return err
	}

	if pos != f.pos {
		return fmt.Errorf("ht: cannot seek to %d in file %s, got to %d", f.pos, f.Path, pos)
	}
	written, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	if err := f.FulfilledBytes(written); err != nil {
		return err
	}
	for _, disallow := range f.Disallow {
		if bytes.Contains(written, []byte(disallow)) {
			return fmt.Errorf("found forbidden %q", disallow)
		}
	}
	return nil
}

// Prepare implements Check's Prepare method.
func (f *Logfile) Prepare() error {
	file, err := os.Open(f.Path)
	if err != nil {
		f.pos = 0
		return nil
	}
	defer file.Close()
	stat, err := file.Stat()
	if err != nil {
		return err
	}

	f.pos = stat.Size()
	return nil
}
<commit_msg>ht: improve documentation<commit_after>// Copyright 2015 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ht

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
)

func init() {
	RegisterCheck(&Logfile{})
}

// ----------------------------------------------------------------------------
// Logfile

// Logfile provides checks on a file (i.e. it ignores the HTTP response).
//
// During preparation the current size of the file identified by Path is
// determined. When the check executes it seeks to that position and
// examines anything written to the file since the preparation of the check.
type Logfile struct {
	// Path is the file system path to the logfile.
	Path string

	// Condition the written stuff must fulfill.
	Condition `json:",omitempty"`

	// Disallow states what is forbidden in the written log.
	Disallow []string `json:",omitempty"`

	pos int64
}

// Execute implements Check's Execute method.
func (f *Logfile) Execute(t *Test) error {
	file, err := os.Open(f.Path)
	if err != nil {
		return err
	}
	pos, err := file.Seek(f.pos, os.SEEK_SET)
	if err != nil {
		return err
	}

	if pos != f.pos {
		return fmt.Errorf("ht: cannot seek to %d in file %s, got to %d", f.pos, f.Path, pos)
	}
	written, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	if err := f.FulfilledBytes(written); err != nil {
		return err
	}
	for _, disallow := range f.Disallow {
		if bytes.Contains(written, []byte(disallow)) {
			return fmt.Errorf("found forbidden %q", disallow)
		}
	}
	return nil
}

// Prepare implements Check's Prepare method.
func (f *Logfile) Prepare() error {
	file, err := os.Open(f.Path)
	if err != nil {
		f.pos = 0
		return nil
	}
	defer file.Close()
	stat, err := file.Stat()
	if err != nil {
		return err
	}

	f.pos = stat.Size()
	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Lawrence Woodman <lwoodman@vlifesystems.com>
 */
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/lawrencewoodman/rulehunter"
	"github.com/lawrencewoodman/rulehunter/csvinput"
	"os"
	"path/filepath"
	"runtime"
)

func processExperiment(experimentFilename string, config *config) error {
	var experiment *rulehunter.Experiment
	var p *os.File
	var err error

	progressFullFilename := filepath.Join(
		config.ProgressDir,
		fmt.Sprintf("%s.progress", experimentFilename),
	)
	p, err = os.Create(progressFullFilename)
	if err != nil {
		msg := fmt.Sprintf("Couldn't create progress file: %s", err)
		err := errors.New(msg)
		return err
	}
	defer p.Close()

	experimentFullFilename :=
		filepath.Join(config.ExperimentsDir, experimentFilename)
	experiment, err = loadExperiment(experimentFullFilename)
	if err != nil {
		msg := fmt.Sprintf("Couldn't load experiment file: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	reportProgress(p, "Describing input")
	fieldDescriptions, err := rulehunter.DescribeInput(experiment.Input)
	if err != nil {
		msg := fmt.Sprintf("Couldn't describe input: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	reportProgress(p, "Generating rules")
	rules, err :=
		rulehunter.GenerateRules(fieldDescriptions, experiment.ExcludeFieldNames)
	if err != nil {
		msg := fmt.Sprintf("Couldn't make report: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment, err := assessRules(rules, experiment, p)
	if err != nil {
		msg := fmt.Sprintf("Couldn't assess rules: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment.Sort(experiment.SortOrder)
	assessment.Refine(3)
	sortedRules := assessment.GetRules()

	reportProgress(p, "Tweaking rules")
	tweakableRules := rulehunter.TweakRules(sortedRules, fieldDescriptions)

	assessment2, err := assessRules(tweakableRules, experiment, p)
	if err != nil {
		msg := fmt.Sprintf("Couldn't assess rules: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment3, err := assessment.Merge(assessment2)
	if err != nil {
		msg := fmt.Sprintf("Couldn't merge assessments: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}
	assessment3.Sort(experiment.SortOrder)
	assessment3.Refine(1)

	bestNonCombinedRules := assessment3.GetRules()
	numRulesToCombine := 50
	if len(bestNonCombinedRules) < numRulesToCombine {
		numRulesToCombine = len(bestNonCombinedRules)
	}
	combinedRules :=
		rulehunter.CombineRules(bestNonCombinedRules[:numRulesToCombine])

	assessment4, err := assessRules(combinedRules, experiment, p)
	if err != nil {
		msg := fmt.Sprintf("Couldn't assess rules: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment5, err := assessment3.Merge(assessment4)
	if err != nil {
		msg := fmt.Sprintf("Couldn't merge assessments: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	err =
		writeReport(assessment5, experiment, experimentFilename, config.ReportsDir)
	if err != nil {
		reportProgress(p, err.Error())
		return err
	}
	return nil
}

type experimentFile struct {
	Title                 string
	InputFilename         string
	FieldNames            []string
	ExcludeFieldNames     []string
	IsFirstLineFieldNames bool
	Separator             string
	Aggregators           []*rulehunter.AggregatorDesc
	Goals                 []string
	SortOrder             []*rulehunter.SortDesc
}

func loadExperiment(filename string) (*rulehunter.Experiment, error) {
	var f *os.File
	var e experimentFile
	var err error

	f, err = os.Open(filename)
	if err != nil {
		return nil, err
	}

	dec := json.NewDecoder(f)
	if err = dec.Decode(&e); err != nil {
		return nil, err
	}

	input, err := csvinput.New(
		e.FieldNames,
		e.InputFilename,
		rune(e.Separator[0]),
		e.IsFirstLineFieldNames,
	)
	if err != nil {
		return nil, err
	}
	experimentDesc := &rulehunter.ExperimentDesc{
		Title:         e.Title,
		Input:         input,
		Fields:        e.FieldNames,
		ExcludeFields: e.ExcludeFieldNames,
		Aggregators:   e.Aggregators,
		Goals:         e.Goals,
		SortOrder:     e.SortOrder,
	}
	experiment, err := rulehunter.MakeExperiment(experimentDesc)
	return experiment, err
}

func reportProgress(f *os.File, msg string) {
	f.WriteString(fmt.Sprintf("%s\n", msg))
}

func prettyPrintFieldDescriptions(fds map[string]*rulehunter.FieldDescription) {
	fmt.Println("Input Description\n")
	for field, fd := range fds {
		fmt.Println("--------------------------")
		fmt.Printf("%s\n--------------------------\n", field)
		prettyPrintFieldDescription(fd)
	}
	fmt.Println("\n")
}

func prettyPrintFieldDescription(fd *rulehunter.FieldDescription) {
	fmt.Printf("Kind: %s\n", fd.Kind)
	fmt.Printf("Min: %s\n", fd.Min)
	fmt.Printf("Max: %s\n", fd.Max)
	fmt.Printf("MaxDP: %d\n", fd.MaxDP)
	fmt.Printf("Values: %s\n", fd.Values)
}

func assessRules(
	rules []*rulehunter.Rule,
	experiment *rulehunter.Experiment,
	progressFile *os.File,
) (*rulehunter.Assessment, error) {
	var assessment *rulehunter.Assessment
	// TODO: Make this part of the config
	maxProcesses := runtime.NumCPU()
	c := make(chan *rulehunter.AssessRulesMPOutcome)

	reportProgress(progressFile,
		fmt.Sprintf("Assessing rules using %d CPUs...\n", maxProcesses))
	go rulehunter.AssessRulesMP(
		rules,
		experiment.Aggregators,
		experiment.Goals,
		experiment.Input,
		maxProcesses,
		c,
	)
	for o := range c {
		if o.Err != nil {
			return nil, o.Err
		}
		reportProgress(progressFile, fmt.Sprintf("Progress: %.2f%%", o.Progress*100))
		assessment = o.Assessment
	}
	reportProgress(progressFile, "Progress: complete")
	return assessment, nil
}

func moveExperimentToSuccess(experimentFilename string, config *config) error {
	experimentFullFilename :=
		filepath.Join(config.ExperimentsDir, experimentFilename)
	experimentSuccessFullFilename :=
		filepath.Join(config.ExperimentsDir, "success", experimentFilename)
	return os.Rename(experimentFullFilename, experimentSuccessFullFilename)
}

func moveExperimentToFail(experimentFilename string, config *config) error {
	experimentFullFilename :=
		filepath.Join(config.ExperimentsDir, experimentFilename)
	experimentFailFullFilename :=
		filepath.Join(config.ExperimentsDir, "fail", experimentFilename)
	return os.Rename(experimentFullFilename, experimentFailFullFilename)
}
<commit_msg>Close an experiment when finished processing it<commit_after>/*
 * Copyright (C) 2016 Lawrence Woodman <lwoodman@vlifesystems.com>
 */
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/lawrencewoodman/rulehunter"
	"github.com/lawrencewoodman/rulehunter/csvinput"
	"os"
	"path/filepath"
	"runtime"
)

func processExperiment(experimentFilename string, config *config) error {
	var experiment *rulehunter.Experiment
	var p *os.File
	var err error

	progressFullFilename := filepath.Join(
		config.ProgressDir,
		fmt.Sprintf("%s.progress", experimentFilename),
	)
	p, err = os.Create(progressFullFilename)
	if err != nil {
		msg := fmt.Sprintf("Couldn't create progress file: %s", err)
		err := errors.New(msg)
		return err
	}
	defer p.Close()

	experimentFullFilename :=
		filepath.Join(config.ExperimentsDir, experimentFilename)
	experiment, err = loadExperiment(experimentFullFilename)
	if err != nil {
		msg := fmt.Sprintf("Couldn't load experiment file: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}
	defer experiment.Close()

	reportProgress(p, "Describing input")
	fieldDescriptions, err := rulehunter.DescribeInput(experiment.Input)
	if err != nil {
		msg := fmt.Sprintf("Couldn't describe input: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	reportProgress(p, "Generating rules")
	rules, err :=
		rulehunter.GenerateRules(fieldDescriptions, experiment.ExcludeFieldNames)
	if err != nil {
		msg := fmt.Sprintf("Couldn't make report: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment, err := assessRules(rules, experiment, p)
	if err != nil {
		msg := fmt.Sprintf("Couldn't assess rules: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment.Sort(experiment.SortOrder)
	assessment.Refine(3)
	sortedRules := assessment.GetRules()

	reportProgress(p, "Tweaking rules")
	tweakableRules := rulehunter.TweakRules(sortedRules, fieldDescriptions)

	assessment2, err := assessRules(tweakableRules, experiment, p)
	if err != nil {
		msg := fmt.Sprintf("Couldn't assess rules: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment3, err := assessment.Merge(assessment2)
	if err != nil {
		msg := fmt.Sprintf("Couldn't merge assessments: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}
	assessment3.Sort(experiment.SortOrder)
	assessment3.Refine(1)

	bestNonCombinedRules := assessment3.GetRules()
	numRulesToCombine := 50
	if len(bestNonCombinedRules) < numRulesToCombine {
		numRulesToCombine = len(bestNonCombinedRules)
	}
	combinedRules :=
		rulehunter.CombineRules(bestNonCombinedRules[:numRulesToCombine])

	assessment4, err := assessRules(combinedRules, experiment, p)
	if err != nil {
		msg := fmt.Sprintf("Couldn't assess rules: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	assessment5, err := assessment3.Merge(assessment4)
	if err != nil {
		msg := fmt.Sprintf("Couldn't merge assessments: %s", err)
		reportProgress(p, msg)
		err := errors.New(msg)
		return err
	}

	err =
		writeReport(assessment5, experiment, experimentFilename, config.ReportsDir)
	if err != nil {
		reportProgress(p, err.Error())
		return err
	}
	return nil
}

type experimentFile struct {
	Title                 string
	InputFilename         string
	FieldNames            []string
	ExcludeFieldNames     []string
	IsFirstLineFieldNames bool
	Separator             string
	Aggregators           []*rulehunter.AggregatorDesc
	Goals                 []string
	SortOrder             []*rulehunter.SortDesc
}

func loadExperiment(filename string) (*rulehunter.Experiment, error) {
	var f *os.File
	var e experimentFile
	var err error

	f, err = os.Open(filename)
	if err != nil {
		return nil, err
	}

	dec := json.NewDecoder(f)
	if err = dec.Decode(&e); err != nil {
		return nil, err
	}

	input, err := csvinput.New(
		e.FieldNames,
		e.InputFilename,
		rune(e.Separator[0]),
		e.IsFirstLineFieldNames,
	)
	if err != nil {
		return nil, err
	}
	experimentDesc := &rulehunter.ExperimentDesc{
		Title:         e.Title,
		Input:         input,
		Fields:        e.FieldNames,
		ExcludeFields: e.ExcludeFieldNames,
		Aggregators:   e.Aggregators,
		Goals:         e.Goals,
		SortOrder:     e.SortOrder,
	}
	experiment, err := rulehunter.MakeExperiment(experimentDesc)
	return experiment, err
}

func reportProgress(f *os.File, msg string) {
	f.WriteString(fmt.Sprintf("%s\n", msg))
}

func prettyPrintFieldDescriptions(fds map[string]*rulehunter.FieldDescription) {
	fmt.Println("Input Description\n")
	for field, fd := range fds {
		fmt.Println("--------------------------")
		fmt.Printf("%s\n--------------------------\n", field)
		prettyPrintFieldDescription(fd)
	}
	fmt.Println("\n")
}

func prettyPrintFieldDescription(fd *rulehunter.FieldDescription) {
	fmt.Printf("Kind: %s\n", fd.Kind)
	fmt.Printf("Min: %s\n", fd.Min)
	fmt.Printf("Max: %s\n", fd.Max)
	fmt.Printf("MaxDP: %d\n", fd.MaxDP)
	fmt.Printf("Values: %s\n", fd.Values)
}

func assessRules(
	rules []*rulehunter.Rule,
	experiment *rulehunter.Experiment,
	progressFile *os.File,
) (*rulehunter.Assessment, error) {
	var assessment *rulehunter.Assessment
	// TODO: Make this part of the config
	maxProcesses := runtime.NumCPU()
	c := make(chan *rulehunter.AssessRulesMPOutcome)

	reportProgress(progressFile,
		fmt.Sprintf("Assessing rules using %d CPUs...\n", maxProcesses))
	go rulehunter.AssessRulesMP(
		rules,
		experiment.Aggregators,
		experiment.Goals,
		experiment.Input,
		maxProcesses,
		c,
	)
	for o := range c {
		if o.Err != nil {
			return nil, o.Err
		}
		reportProgress(progressFile, fmt.Sprintf("Progress: %.2f%%", o.Progress*100))
		assessment = o.Assessment
	}
	reportProgress(progressFile, "Progress: complete")
	return assessment, nil
}

func moveExperimentToSuccess(experimentFilename string, config *config) error {
	experimentFullFilename :=
		filepath.Join(config.ExperimentsDir, experimentFilename)
	experimentSuccessFullFilename :=
		filepath.Join(config.ExperimentsDir, "success", experimentFilename)
	return os.Rename(experimentFullFilename, experimentSuccessFullFilename)
}

func moveExperimentToFail(experimentFilename string, config *config) error {
	experimentFullFilename :=
		filepath.Join(config.ExperimentsDir, experimentFilename)
	experimentFailFullFilename :=
		filepath.Join(config.ExperimentsDir, "fail", experimentFilename)
	return os.Rename(experimentFullFilename, experimentFailFullFilename)
}
<|endoftext|>