<commit_before>package mail

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/textproto"
	"os"
	"regexp"
	"strings"
	"testing"
)

func parseMailFromString(source string) (*MimeEntity, error) {
	return parseMailFromStringWithGpg(source, nil)
}

func parseMailFromStringWithGpg(source string, gpg GpgUtility) (*MimeEntity, error) {
	parser := Parser{gpg}
	reader := strings.NewReader(source)
	return parser.ParseMail(reader)
}

func loadTestMail(fileName string) ([]byte, error) {
	input, err := os.Open("../test/mails/" + fileName)
	if err != nil {
		return nil, err
	}
	defer func() { input.Close() }()
	data, err := ioutil.ReadAll(input)
	if err != nil {
		return nil, err
	}
	regex := regexp.MustCompile("\r?\n")
	return regex.ReplaceAll(data, []byte("\r\n")), nil
}

func parseMailFromFile(fileName string) (*MimeEntity, error) {
	return parseMailFromFileWithGpg(fileName, nil)
}

func parseMailFromFileWithGpg(fileName string, gpg GpgUtility) (*MimeEntity, error) {
	data, err := loadTestMail(fileName)
	if err != nil {
		return nil, err
	}
	parser := Parser{gpg}
	return parser.ParseMail(bytes.NewReader(data))
}

type MultipartBuilder struct {
	Boundary string
	Buffer   bytes.Buffer
}

func createMultipart(boundary string) *MultipartBuilder {
	return &MultipartBuilder{Boundary: boundary}
}

func (builder *MultipartBuilder) withPart(contentType string,
	text string, header textproto.MIMEHeader) *MultipartBuilder {
	builder.Buffer.WriteString(fmt.Sprintf("--%s\r\n", builder.Boundary))
	if contentType != "" {
		builder.Buffer.WriteString(fmt.Sprintf("Content-Type: %s\r\n", contentType))
	}
	if header != nil {
		for key, values := range header {
			for _, value := range values {
				builder.Buffer.WriteString(fmt.Sprintf("%s: %s\r\n", key, value))
			}
		}
	}
	builder.Buffer.WriteString("\r\n")
	builder.Buffer.WriteString(text)
	builder.Buffer.WriteString("\r\n")
	return builder
}

func (builder *MultipartBuilder) build() string {
	builder.Buffer.WriteString(fmt.Sprintf("--%s--\r\n", builder.Boundary))
	return builder.Buffer.String()
}

func TestGetMimeMediaTypeFromHeader(t *testing.T) {
	header := textproto.MIMEHeader{"Content-Type": {"foo; param=bar", "stuff"}}
	mediaType, err := getMimeMediaTypeFromHeader(header, "Content-Type", "")
	assert.NoError(t, err)
	assert.Equal(t, "foo", mediaType.Value)
	assert.Equal(t, 1, len(mediaType.Params))
	assert.Equal(t, "bar", mediaType.Params["param"])
	mediaType, err = getMimeMediaTypeFromHeader(header, "key", "default")
	assert.NoError(t, err)
	assert.Equal(t, "default", mediaType.Value)
}

func TestParseHeaders(t *testing.T) {
	mail, err := parseMailFromFile("plaintext.eml")
	assert.NoError(t, err)
	assert.Equal(t, 8, len(mail.Header))
	assert.Equal(t, "Basic test", mail.Header["Subject"][0])
	assert.Equal(t, "text/plain", mail.Header["Content-Type"][0])
}

func TestParseEmptyMail(t *testing.T) {
	mail, err := parseMailFromString("\r\n")
	assert.NoError(t, err)
	assert.Equal(t, 0, len(mail.Parts))
	assert.Equal(t, 0, len(mail.Header))
	assert.False(t, mail.IsAttachment)
	assert.False(t, mail.IsSigned)
}

func TestParseText(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"text/plain", nil}
	header := textproto.MIMEHeader{}
	entity, err := parser.parseText(contentType, header, strings.NewReader("Hello world!"))
	assert.NoError(t, err)
	assert.Equal(t, "Hello world!", string(entity.Content))
}

func TestParseTextWithCharset(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"text/plain", map[string]string{"charset": "ISO-8859-1"}}
	header := textproto.MIMEHeader{}
	entity, err := parser.parseText(contentType, header, strings.NewReader("\xe0\x20\x63\xf4\x74\xe9"))
	assert.NoError(t, err)
	assert.Equal(t, "à côté", string(entity.Content))
}

func TestParseTextWithInvalidEncoding(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"text/plain", map[string]string{"charset": "utf-8"}}
	header := textproto.MIMEHeader{}
	entity, err := parser.parseText(contentType, header, bytes.NewReader([]byte{0xC0}))
	assert.Nil(t, err)
	assert.Equal(t, "\uFFFD", string(entity.Content))
}

func TestParseMultipart(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	text := createMultipart("frontier").
		withPart("text/plain", "Part0", nil).
		withPart("text/plain", "Part1", nil).
		withPart("text/plain", "Part2", nil).
		build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NoError(t, err)
	assert.Equal(t, 3, len(entity.Parts))
	assert.Equal(t, "Part0", string(entity.Parts[0].Content))
	assert.Equal(t, "Part1", string(entity.Parts[1].Content))
	assert.Equal(t, "Part2", string(entity.Parts[2].Content))
}

func TestParseMultipartWithoutBoundary(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{}}
	header := textproto.MIMEHeader{}
	text := createMultipart("frontier").withPart("text/plain", "Hello world!", nil).build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NotNil(t, err)
	assert.Nil(t, entity)
}

func TestParseMultipartNested(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	inner := createMultipart("inner").withPart("text/plain", "Nested text", nil).build()
	text := createMultipart("frontier").
		withPart("multipart/mixed; boundary=inner", inner, nil).
		build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NoError(t, err)
	assert.Equal(t, 1, len(entity.Parts))
	innerPart := entity.Parts[0]
	assert.Equal(t, 1, len(innerPart.Header))
	assert.Equal(t, "multipart/mixed; boundary=inner", innerPart.Header.Get("Content-Type"))
	assert.Equal(t, 1, len(innerPart.Parts))
	assert.Equal(t, "Nested text", string(innerPart.Parts[0].Content))
}

func TestParseMultipartInvalid(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	text := createMultipart("frontier").
		withPart("text/plain", "Hello", nil).
		build()
	text = text[:len(text)-4] // Remove final boundary end '--' (and \r\n).
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.Nil(t, entity)
	assert.NotNil(t, err)
}

func TestParseMultipartWithQuotedPrintable(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	innerHeader := textproto.MIMEHeader{}
	innerHeader.Set("Content-Type", "text/plain; charset=ISO-8859-1")
	innerHeader.Set("Content-Transfer-Encoding", "quoted-printable")
	text := createMultipart("frontier").withPart("", "B=E4renf=FC=DFe", innerHeader).build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NoError(t, err)
	assert.Equal(t, "Bärenfüße", string(entity.Parts[0].Content))
}

type MockGpg struct {
	t                                                     *testing.T
	expectedMicAlgorithm, expectedData, expectedSignature string
	checked                                               bool
}

func (gpg *MockGpg) CheckSignature(micAlgorithm string, data []byte, signature []byte) bool {
	assert.Equal(gpg.t, gpg.expectedMicAlgorithm, micAlgorithm)
	assert.Equal(gpg.t, gpg.expectedData, string(data))
	assert.Equal(gpg.t, gpg.expectedSignature, string(signature))
	gpg.checked = true
	return true
}

func TestParseMultipartSigned(t *testing.T) {
	text := createMultipart("frontier").
		withPart("text/plain", "Hello there!", nil).
		withPart("application/pgp-signature", "SIGNATURE", nil).
		build()

	expectedSignedPart := "Content-Type: text/plain\r\n\r\nHello there!\r\n"
	mockGpg := &MockGpg{t, "pgp-sha1", expectedSignedPart, "SIGNATURE", false}
	parser := Parser{mockGpg}
	contentType := MimeMediaType{"multipart/signed", map[string]string{"boundary": "frontier", "micalg": "pgp-sha1"}}
	mail, err := parser.parseMultipartSigned(contentType, textproto.MIMEHeader{}, strings.NewReader(text))
	assert.NoError(t, err)
	assert.True(t, mail.IsSigned)
	assert.True(t, mockGpg.checked)
	assert.Equal(t, 2, len(mail.Parts))
}

func TestPlainText(t *testing.T) {
	mail, err := parseMailFromFile("plaintext.eml")
	assert.NoError(t, err)
	assert.Equal(t, "This is some nice plain text!\r\n", string(mail.Content))
}

func TestMultipartSigned(t *testing.T) {
	signedPart := "Content-Type: text/plain\r\n\r\nHello this is me!\r\n\r\n"
	mockGpg := &MockGpg{t, "pgp-sha1", signedPart, "SIGNATURE", false}
	mail, err := parseMailFromFileWithGpg("signed_multipart_simple.eml", mockGpg)
	assert.NoError(t, err)
	assert.True(t, mockGpg.checked)
	assert.Equal(t, 2, len(mail.Parts))
	assert.True(t, mail.IsSigned)
}

func TestFindAttachment(t *testing.T) {
	mail, err := parseMailFromFile("attachment.eml")
	assert.NoError(t, err)
	data := mail.FindAttachment("application/octet-stream")
	assert.Equal(t, "This is a PDF file.", string(data))
	data = mail.FindAttachment("image/png")
	assert.Nil(t, data)
}
<commit_msg>Added a comment.<commit_after>package mail

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/textproto"
	"os"
	"regexp"
	"strings"
	"testing"
)

func parseMailFromString(source string) (*MimeEntity, error) {
	return parseMailFromStringWithGpg(source, nil)
}

func parseMailFromStringWithGpg(source string, gpg GpgUtility) (*MimeEntity, error) {
	parser := Parser{gpg}
	reader := strings.NewReader(source)
	return parser.ParseMail(reader)
}

func loadTestMail(fileName string) ([]byte, error) {
	input, err := os.Open("../test/mails/" + fileName)
	if err != nil {
		return nil, err
	}
	defer func() { input.Close() }()
	data, err := ioutil.ReadAll(input)
	if err != nil {
		return nil, err
	}
	// Do not rely on correct line endings in test files, ensure them here.
	regex := regexp.MustCompile("\r?\n")
	return regex.ReplaceAll(data, []byte("\r\n")), nil
}

func parseMailFromFile(fileName string) (*MimeEntity, error) {
	return parseMailFromFileWithGpg(fileName, nil)
}

func parseMailFromFileWithGpg(fileName string, gpg GpgUtility) (*MimeEntity, error) {
	data, err := loadTestMail(fileName)
	if err != nil {
		return nil, err
	}
	parser := Parser{gpg}
	return parser.ParseMail(bytes.NewReader(data))
}

type MultipartBuilder struct {
	Boundary string
	Buffer   bytes.Buffer
}

func createMultipart(boundary string) *MultipartBuilder {
	return &MultipartBuilder{Boundary: boundary}
}

func (builder *MultipartBuilder) withPart(contentType string,
	text string, header textproto.MIMEHeader) *MultipartBuilder {
	builder.Buffer.WriteString(fmt.Sprintf("--%s\r\n", builder.Boundary))
	if contentType != "" {
		builder.Buffer.WriteString(fmt.Sprintf("Content-Type: %s\r\n", contentType))
	}
	if header != nil {
		for key, values := range header {
			for _, value := range values {
				builder.Buffer.WriteString(fmt.Sprintf("%s: %s\r\n", key, value))
			}
		}
	}
	builder.Buffer.WriteString("\r\n")
	builder.Buffer.WriteString(text)
	builder.Buffer.WriteString("\r\n")
	return builder
}

func (builder *MultipartBuilder) build() string {
	builder.Buffer.WriteString(fmt.Sprintf("--%s--\r\n", builder.Boundary))
	return builder.Buffer.String()
}

func TestGetMimeMediaTypeFromHeader(t *testing.T) {
	header := textproto.MIMEHeader{"Content-Type": {"foo; param=bar", "stuff"}}
	mediaType, err := getMimeMediaTypeFromHeader(header, "Content-Type", "")
	assert.NoError(t, err)
	assert.Equal(t, "foo", mediaType.Value)
	assert.Equal(t, 1, len(mediaType.Params))
	assert.Equal(t, "bar", mediaType.Params["param"])
	mediaType, err = getMimeMediaTypeFromHeader(header, "key", "default")
	assert.NoError(t, err)
	assert.Equal(t, "default", mediaType.Value)
}

func TestParseHeaders(t *testing.T) {
	mail, err := parseMailFromFile("plaintext.eml")
	assert.NoError(t, err)
	assert.Equal(t, 8, len(mail.Header))
	assert.Equal(t, "Basic test", mail.Header["Subject"][0])
	assert.Equal(t, "text/plain", mail.Header["Content-Type"][0])
}

func TestParseEmptyMail(t *testing.T) {
	mail, err := parseMailFromString("\r\n")
	assert.NoError(t, err)
	assert.Equal(t, 0, len(mail.Parts))
	assert.Equal(t, 0, len(mail.Header))
	assert.False(t, mail.IsAttachment)
	assert.False(t, mail.IsSigned)
}

func TestParseText(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"text/plain", nil}
	header := textproto.MIMEHeader{}
	entity, err := parser.parseText(contentType, header, strings.NewReader("Hello world!"))
	assert.NoError(t, err)
	assert.Equal(t, "Hello world!", string(entity.Content))
}

func TestParseTextWithCharset(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"text/plain", map[string]string{"charset": "ISO-8859-1"}}
	header := textproto.MIMEHeader{}
	entity, err := parser.parseText(contentType, header, strings.NewReader("\xe0\x20\x63\xf4\x74\xe9"))
	assert.NoError(t, err)
	assert.Equal(t, "à côté", string(entity.Content))
}

func TestParseTextWithInvalidEncoding(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"text/plain", map[string]string{"charset": "utf-8"}}
	header := textproto.MIMEHeader{}
	entity, err := parser.parseText(contentType, header, bytes.NewReader([]byte{0xC0}))
	assert.Nil(t, err)
	assert.Equal(t, "\uFFFD", string(entity.Content))
}

func TestParseMultipart(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	text := createMultipart("frontier").
		withPart("text/plain", "Part0", nil).
		withPart("text/plain", "Part1", nil).
		withPart("text/plain", "Part2", nil).
		build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NoError(t, err)
	assert.Equal(t, 3, len(entity.Parts))
	assert.Equal(t, "Part0", string(entity.Parts[0].Content))
	assert.Equal(t, "Part1", string(entity.Parts[1].Content))
	assert.Equal(t, "Part2", string(entity.Parts[2].Content))
}

func TestParseMultipartWithoutBoundary(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{}}
	header := textproto.MIMEHeader{}
	text := createMultipart("frontier").withPart("text/plain", "Hello world!", nil).build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NotNil(t, err)
	assert.Nil(t, entity)
}

func TestParseMultipartNested(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	inner := createMultipart("inner").withPart("text/plain", "Nested text", nil).build()
	text := createMultipart("frontier").
		withPart("multipart/mixed; boundary=inner", inner, nil).
		build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NoError(t, err)
	assert.Equal(t, 1, len(entity.Parts))
	innerPart := entity.Parts[0]
	assert.Equal(t, 1, len(innerPart.Header))
	assert.Equal(t, "multipart/mixed; boundary=inner", innerPart.Header.Get("Content-Type"))
	assert.Equal(t, 1, len(innerPart.Parts))
	assert.Equal(t, "Nested text", string(innerPart.Parts[0].Content))
}

func TestParseMultipartInvalid(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	text := createMultipart("frontier").
		withPart("text/plain", "Hello", nil).
		build()
	text = text[:len(text)-4] // Remove final boundary end '--' (and \r\n).
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.Nil(t, entity)
	assert.NotNil(t, err)
}

func TestParseMultipartWithQuotedPrintable(t *testing.T) {
	parser := &Parser{Gpg: nil}
	contentType := MimeMediaType{"multipart/mixed", map[string]string{"boundary": "frontier"}}
	header := textproto.MIMEHeader{}
	innerHeader := textproto.MIMEHeader{}
	innerHeader.Set("Content-Type", "text/plain; charset=ISO-8859-1")
	innerHeader.Set("Content-Transfer-Encoding", "quoted-printable")
	text := createMultipart("frontier").withPart("", "B=E4renf=FC=DFe", innerHeader).build()
	entity, err := parser.parseMultipart(contentType, header, strings.NewReader(text))
	assert.NoError(t, err)
	assert.Equal(t, "Bärenfüße", string(entity.Parts[0].Content))
}

type MockGpg struct {
	t                                                     *testing.T
	expectedMicAlgorithm, expectedData, expectedSignature string
	checked                                               bool
}

func (gpg *MockGpg) CheckSignature(micAlgorithm string, data []byte, signature []byte) bool {
	assert.Equal(gpg.t, gpg.expectedMicAlgorithm, micAlgorithm)
	assert.Equal(gpg.t, gpg.expectedData, string(data))
	assert.Equal(gpg.t, gpg.expectedSignature, string(signature))
	gpg.checked = true
	return true
}

func TestParseMultipartSigned(t *testing.T) {
	text := createMultipart("frontier").
		withPart("text/plain", "Hello there!", nil).
		withPart("application/pgp-signature", "SIGNATURE", nil).
		build()

	expectedSignedPart := "Content-Type: text/plain\r\n\r\nHello there!\r\n"
	mockGpg := &MockGpg{t, "pgp-sha1", expectedSignedPart, "SIGNATURE", false}
	parser := Parser{mockGpg}
	contentType := MimeMediaType{"multipart/signed", map[string]string{"boundary": "frontier", "micalg": "pgp-sha1"}}
	mail, err := parser.parseMultipartSigned(contentType, textproto.MIMEHeader{}, strings.NewReader(text))
	assert.NoError(t, err)
	assert.True(t, mail.IsSigned)
	assert.True(t, mockGpg.checked)
	assert.Equal(t, 2, len(mail.Parts))
}

func TestPlainText(t *testing.T) {
	mail, err := parseMailFromFile("plaintext.eml")
	assert.NoError(t, err)
	assert.Equal(t, "This is some nice plain text!\r\n", string(mail.Content))
}

func TestMultipartSigned(t *testing.T) {
	signedPart := "Content-Type: text/plain\r\n\r\nHello this is me!\r\n\r\n"
	mockGpg := &MockGpg{t, "pgp-sha1", signedPart, "SIGNATURE", false}
	mail, err := parseMailFromFileWithGpg("signed_multipart_simple.eml", mockGpg)
	assert.NoError(t, err)
	assert.True(t, mockGpg.checked)
	assert.Equal(t, 2, len(mail.Parts))
	assert.True(t, mail.IsSigned)
}

func TestFindAttachment(t *testing.T) {
	mail, err := parseMailFromFile("attachment.eml")
	assert.NoError(t, err)
	data := mail.FindAttachment("application/octet-stream")
	assert.Equal(t, "This is a PDF file.", string(data))
	data = mail.FindAttachment("image/png")
	assert.Nil(t, data)
}
<|endoftext|>


<commit_before>package sgload

import (
	"fmt"
	"strings"
	"sync"
	"time"

	sgreplicate "github.com/couchbaselabs/sg-replicate"
	"github.com/peterbourgon/g2s"
)

type Reader struct {
	Agent
	SGChannels      []string    // The Sync Gateway channels this reader is assigned to pull from
	NumDocsExpected int         // The total number of docs this reader is expected to pull
	BatchSize       int         // The number of docs to pull in batch (_changes feed and bulk_get)
	StatsdClient    *g2s.Statsd // The statsd client instance to use to push stats to statdsd
}

func NewReader(wg *sync.WaitGroup, ID int, u UserCred, d DataStore, batchsize int) *Reader {

	return &Reader{
		Agent: Agent{
			FinishedWg: wg,
			UserCred:   u,
			ID:         ID,
			DataStore:  d,
			BatchSize:  batchsize,
		},
	}
}

func (r *Reader) SetChannels(sgChannels []string) {
	r.SGChannels = sgChannels
}

func (r *Reader) SetNumDocsExpected(n int) {
	r.NumDocsExpected = n
}

func (r *Reader) SetBatchSize(batchSize int) {
	r.BatchSize = batchSize
}

func (r *Reader) SetStatsdClient(statsdClient *g2s.Statsd) {
	r.StatsdClient = statsdClient
}

func (r *Reader) Run() {

	numDocsPulled := 0
	since := StringSincer{}
	result := pullMoreDocsResult{}
	var err error
	var timeStartedCreatingDocs time.Time

	defer r.FinishedWg.Done()
	defer func() {
		delta := time.Since(timeStartedCreatingDocs)
		r.StatsdClient.Timing(
			statsdSampleRate,
			"get_all_documents",
			delta,
		)
	}()

	r.createSGUserIfNeeded()

	timeStartedCreatingDocs = time.Now()

	for {

		if r.isFinished(numDocsPulled) {
			logger.Info("Reader finished", "agent.ID", r.ID, "numDocs", r.NumDocsExpected)
			break
		}
		result, err = r.pullMoreDocs(since)
		if err != nil {
			panic(fmt.Sprintf("Got error getting changes: %v", err))
		}
		since = result.since
		numDocsPulled += result.numDocsPulled

	}

}

func (r *Reader) createSGUserIfNeeded() {
	if r.CreateDataStoreUser == true {

		logger.Info("Creating reader SG user", "username", r.UserCred.Username, "channels", r.SGChannels)

		if err := r.DataStore.CreateUser(r.UserCred, r.SGChannels); err != nil {
			panic(fmt.Sprintf("Error creating user in datastore.  User: %v, Err: %v", r.UserCred, err))
		}
	}

}

func (r *Reader) isFinished(numDocsPulled int) bool {

	logger.Info("Reader checking if finished", "agent.ID", r.ID, "numDocsPulled", numDocsPulled, "numDocsExpected", r.NumDocsExpected)

	switch {
	case numDocsPulled > r.NumDocsExpected:
		panic(fmt.Sprintf("Reader was only expected to pull %d docs, but pulled %d.", r.NumDocsExpected, numDocsPulled))
	case numDocsPulled == r.NumDocsExpected:
		return true
	default:
		return false
	}

}

type pullMoreDocsResult struct {
	since         StringSincer
	numDocsPulled int
}

func (r *Reader) pullMoreDocs(since Sincer) (pullMoreDocsResult, error) {

	// Create a retry sleeper which controls how many times to retry
	// and how long to wait in between retries
	numRetries := 7
	sleepMsBetweenRetry := 500
	retrySleeper := CreateDoublingSleeperFunc(numRetries, sleepMsBetweenRetry)

	// Create retry worker that knows how to do actual work
	retryWorker := func() (shouldRetry bool, err error, value interface{}) {

		result := pullMoreDocsResult{}

		changes, newSince, err := r.DataStore.Changes(since, r.BatchSize)
		if err != nil {
			return false, err, result
		}

		if len(changes.Results) == 0 {
			logger.Warn("No changes pulled", "agent.ID", r.ID, "since", since)
			return true, nil, result
		}
		if newSince.Equals(since) {
			logger.Warn("Since value should have changed", "agent.ID", r.ID, "since", since, "newsince", newSince)
			return true, nil, result
		}

		// Strip out any changes with id "id":"_user/*"
		// since they are user docs and we don't care about them
		changes = stripUserDocChanges(changes)

		bulkGetRequest := getBulkGetRequest(changes)

		err = r.DataStore.BulkGetDocuments(bulkGetRequest)
		if err != nil {
			return false, err, result
		}

		result.since = newSince.(StringSincer)
		result.numDocsPulled = len(bulkGetRequest.Docs)
		return false, nil, result

	}

	// Invoke the retry worker / sleeper combo in a loop
	err, workerReturnVal := RetryLoop("pullMoreDocs", retryWorker, retrySleeper)
	if err != nil {
		return pullMoreDocsResult{}, err
	}

	return workerReturnVal.(pullMoreDocsResult), nil

}

func getBulkGetRequest(changes sgreplicate.Changes) sgreplicate.BulkGetRequest {

	bulkDocsRequest := sgreplicate.BulkGetRequest{}
	docs := []sgreplicate.DocumentRevisionPair{}
	for _, change := range changes.Results {
		docRevPair := sgreplicate.DocumentRevisionPair{}
		docRevPair.Id = change.Id
		docRevPair.Revision = change.ChangedRevs[0].Revision
		docs = append(docs, docRevPair)
	}
	bulkDocsRequest.Docs = docs
	return bulkDocsRequest

}

func stripUserDocChanges(changes sgreplicate.Changes) (changesStripped sgreplicate.Changes) {
	changesStripped.LastSequence = changes.LastSequence

	for _, change := range changes.Results {
		if strings.Contains(change.Id, "_user") {
			continue
		}
		changesStripped.Results = append(changesStripped.Results, change)

	}

	return changesStripped
}

// A retry sleeper is called back by the retry loop and passed
// the current retryCount, and should return the amount of milliseconds
// that the retry should sleep.
type RetrySleeper func(retryCount int) (shouldContinue bool, timeTosleepMs int)

// A RetryWorker encapsulates the work being done in a Retry Loop.  The shouldRetry
// return value determines whether the worker will retry, regardless of the err value.
// If the worker has exceeded it's retry attempts, then it will not be called again
// even if it returns shouldRetry = true.
type RetryWorker func() (shouldRetry bool, err error, value interface{})

func RetryLoop(description string, worker RetryWorker, sleeper RetrySleeper) (error, interface{}) {

	numAttempts := 1

	for {
		shouldRetry, err, value := worker()
		if !shouldRetry {
			if err != nil {
				return err, nil
			}
			return nil, value
		}
		shouldContinue, sleepMs := sleeper(numAttempts)
		if !shouldContinue {
			if err == nil {
				err = fmt.Errorf("RetryLoop for %v giving up after %v attempts", description, numAttempts)
			}
			logger.Warn("RetryLoop giving up", "description", description, "numAttempts", numAttempts)
			return err, value
		}
		logger.Warn("RetryLoop will retry soon", "description", description, "sleepMs", sleepMs)

		<-time.After(time.Millisecond * time.Duration(sleepMs))

		numAttempts += 1

	}
}

// Create a RetrySleeper that will double the retry time on every iteration and
// use the given parameters
func CreateDoublingSleeperFunc(maxNumAttempts, initialTimeToSleepMs int) RetrySleeper {

	timeToSleepMs := initialTimeToSleepMs

	sleeper := func(numAttempts int) (bool, int) {
		if numAttempts > maxNumAttempts {
			return false, -1
		}
		if numAttempts > 1 {
			timeToSleepMs *= 2
		}
		return true, timeToSleepMs
	}
	return sleeper

}
<commit_msg>Add get_change_and_document<commit_after>package sgload

import (
	"fmt"
	"strings"
	"sync"
	"time"

	sgreplicate "github.com/couchbaselabs/sg-replicate"
	"github.com/peterbourgon/g2s"
)

type Reader struct {
	Agent
	SGChannels      []string    // The Sync Gateway channels this reader is assigned to pull from
	NumDocsExpected int         // The total number of docs this reader is expected to pull
	BatchSize       int         // The number of docs to pull in batch (_changes feed and bulk_get)
	StatsdClient    *g2s.Statsd // The statsd client instance to use to push stats to statdsd
}

func NewReader(wg *sync.WaitGroup, ID int, u UserCred, d DataStore, batchsize int) *Reader {

	return &Reader{
		Agent: Agent{
			FinishedWg: wg,
			UserCred:   u,
			ID:         ID,
			DataStore:  d,
			BatchSize:  batchsize,
		},
	}
}

func (r *Reader) SetChannels(sgChannels []string) {
	r.SGChannels = sgChannels
}

func (r *Reader) SetNumDocsExpected(n int) {
	r.NumDocsExpected = n
}

func (r *Reader) SetBatchSize(batchSize int) {
	r.BatchSize = batchSize
}

func (r *Reader) SetStatsdClient(statsdClient *g2s.Statsd) {
	r.StatsdClient = statsdClient
}

func (r *Reader) Run() {

	numDocsPulled := 0
	since := StringSincer{}
	result := pullMoreDocsResult{}
	var err error
	var timeStartedCreatingDocs time.Time

	defer r.FinishedWg.Done()
	defer func() {
		// Calculate how long it took for this reader to read
		// all of its docs
		delta := time.Since(timeStartedCreatingDocs)
		r.StatsdClient.Timing(
			statsdSampleRate,
			"get_all_documents",
			delta,
		)
	}()
	defer func() {
		// Calculate the average time it took to read each doc from
		// the changes feed and the doc itself
		delta := time.Since(timeStartedCreatingDocs)
		delta = time.Duration(int64(delta) / int64(numDocsPulled))
		r.StatsdClient.Timing(
			statsdSampleRate,
			"get_change_and_document",
			delta,
		)
	}()

	r.createSGUserIfNeeded()

	timeStartedCreatingDocs = time.Now()

	for {

		if r.isFinished(numDocsPulled) {
			logger.Info("Reader finished", "agent.ID", r.ID, "numDocs", r.NumDocsExpected)
			break
		}
		result, err = r.pullMoreDocs(since)
		if err != nil {
			panic(fmt.Sprintf("Got error getting changes: %v", err))
		}
		since = result.since
		numDocsPulled += result.numDocsPulled

	}

}

func (r *Reader) createSGUserIfNeeded() {
	if r.CreateDataStoreUser == true {

		logger.Info("Creating reader SG user", "username", r.UserCred.Username, "channels", r.SGChannels)

		if err := r.DataStore.CreateUser(r.UserCred, r.SGChannels); err != nil {
			panic(fmt.Sprintf("Error creating user in datastore.  User: %v, Err: %v", r.UserCred, err))
		}
	}

}

func (r *Reader) isFinished(numDocsPulled int) bool {

	logger.Info("Reader checking if finished", "agent.ID", r.ID, "numDocsPulled", numDocsPulled, "numDocsExpected", r.NumDocsExpected)

	switch {
	case numDocsPulled > r.NumDocsExpected:
		panic(fmt.Sprintf("Reader was only expected to pull %d docs, but pulled %d.", r.NumDocsExpected, numDocsPulled))
	case numDocsPulled == r.NumDocsExpected:
		return true
	default:
		return false
	}

}

type pullMoreDocsResult struct {
	since         StringSincer
	numDocsPulled int
}

func (r *Reader) pullMoreDocs(since Sincer) (pullMoreDocsResult, error) {

	// Create a retry sleeper which controls how many times to retry
	// and how long to wait in between retries
	numRetries := 7
	sleepMsBetweenRetry := 500
	retrySleeper := CreateDoublingSleeperFunc(numRetries, sleepMsBetweenRetry)

	// Create retry worker that knows how to do actual work
	retryWorker := func() (shouldRetry bool, err error, value interface{}) {

		result := pullMoreDocsResult{}

		changes, newSince, err := r.DataStore.Changes(since, r.BatchSize)
		if err != nil {
			return false, err, result
		}

		if len(changes.Results) == 0 {
			logger.Warn("No changes pulled", "agent.ID", r.ID, "since", since)
			return true, nil, result
		}
		if newSince.Equals(since) {
			logger.Warn("Since value should have changed", "agent.ID", r.ID, "since", since, "newsince", newSince)
			return true, nil, result
		}

		// Strip out any changes with id "id":"_user/*"
		// since they are user docs and we don't care about them
		changes = stripUserDocChanges(changes)

		bulkGetRequest := getBulkGetRequest(changes)

		err = r.DataStore.BulkGetDocuments(bulkGetRequest)
		if err != nil {
			return false, err, result
		}

		result.since = newSince.(StringSincer)
		result.numDocsPulled = len(bulkGetRequest.Docs)
		return false, nil, result

	}

	// Invoke the retry worker / sleeper combo in a loop
	err, workerReturnVal := RetryLoop("pullMoreDocs", retryWorker, retrySleeper)
	if err != nil {
		return pullMoreDocsResult{}, err
	}

	return workerReturnVal.(pullMoreDocsResult), nil

}

func getBulkGetRequest(changes sgreplicate.Changes) sgreplicate.BulkGetRequest {

	bulkDocsRequest := sgreplicate.BulkGetRequest{}
	docs := []sgreplicate.DocumentRevisionPair{}
	for _, change := range changes.Results {
		docRevPair := sgreplicate.DocumentRevisionPair{}
		docRevPair.Id = change.Id
		docRevPair.Revision = change.ChangedRevs[0].Revision
		docs = append(docs, docRevPair)
	}
	bulkDocsRequest.Docs = docs
	return bulkDocsRequest

}

func stripUserDocChanges(changes sgreplicate.Changes) (changesStripped sgreplicate.Changes) {
	changesStripped.LastSequence = changes.LastSequence

	for _, change := range changes.Results {
		if strings.Contains(change.Id, "_user") {
			continue
		}
		changesStripped.Results = append(changesStripped.Results, change)

	}

	return changesStripped
}

// A retry sleeper is called back by the retry loop and passed
// the current retryCount, and should return the amount of milliseconds
// that the retry should sleep.
type RetrySleeper func(retryCount int) (shouldContinue bool, timeTosleepMs int)

// A RetryWorker encapsulates the work being done in a Retry Loop.  The shouldRetry
// return value determines whether the worker will retry, regardless of the err value.
// If the worker has exceeded it's retry attempts, then it will not be called again
// even if it returns shouldRetry = true.
type RetryWorker func() (shouldRetry bool, err error, value interface{})

func RetryLoop(description string, worker RetryWorker, sleeper RetrySleeper) (error, interface{}) {

	numAttempts := 1

	for {
		shouldRetry, err, value := worker()
		if !shouldRetry {
			if err != nil {
				return err, nil
			}
			return nil, value
		}
		shouldContinue, sleepMs := sleeper(numAttempts)
		if !shouldContinue {
			if err == nil {
				err = fmt.Errorf("RetryLoop for %v giving up after %v attempts", description, numAttempts)
			}
			logger.Warn("RetryLoop giving up", "description", description, "numAttempts", numAttempts)
			return err, value
		}
		logger.Warn("RetryLoop will retry soon", "description", description, "sleepMs", sleepMs)

		<-time.After(time.Millisecond * time.Duration(sleepMs))

		numAttempts += 1

	}
}

// Create a RetrySleeper that will double the retry time on every iteration and
// use the given parameters
func CreateDoublingSleeperFunc(maxNumAttempts, initialTimeToSleepMs int) RetrySleeper {

	timeToSleepMs := initialTimeToSleepMs

	sleeper := func(numAttempts int) (bool, int) {
		if numAttempts > maxNumAttempts {
			return false, -1
		}
		if numAttempts > 1 {
			timeToSleepMs *= 2
		}
		return true, timeToSleepMs
	}
	return sleeper

}
<|endoftext|>


<commit_before>package manifest

import (
	"fmt"

	"github.com/bpicode/fritzctl/console"
	"github.com/bpicode/fritzctl/fritz"
)

// Applier defines the interface to apply a plan to the AHA system.
type Applier interface {
	// Apply performs the changes necessary to transition from src to target configuration. If the target plan
	// could be realized, it returns an error. If the plan was applied successfully, it returns null.
	Apply(src, target *Plan) error
}

// DryRunner is an Applier that only plans changes to the AHA system.
func DryRunner() Applier {
	return &dryRunner{}
}

type dryRunner struct {
}

// Apply does only log the proposed changes.
func (d *dryRunner) Apply(src, target *Plan) error {
	planner := TargetBasedPlanner(justLogSwitchState, justLogThermostat)
	actions, err := planner.Plan(src, target)
	if err != nil {
		return err
	}
	fmt.Println("\n\nThe following actions would be applied by the manifest:\n")
	for _, action := range actions {
		action.Perform(nil)
	}
	return nil
}

func justLogThermostat(before, after Thermostat) Action {
	return &justLogThermostatAction{before: before, after: after}
}

type justLogThermostatAction struct {
	before Thermostat
	after  Thermostat
}

// Perform only logs changes.
func (a *justLogThermostatAction) Perform(f fritz.HomeAutomationApi) error {
	if a.before.Temperature != a.after.Temperature {
		fmt.Printf("\t'%s'\t%.1f°C\t⟶\t%.1f°C\n", a.before.Name, a.before.Temperature, a.after.Temperature)
	}
	return nil
}

type justLogSwitchAction struct {
	before Switch
	after  Switch
}

func justLogSwitchState(before, after Switch) Action {
	return &justLogSwitchAction{before: before, after: after}
}

// Perform only logs changes.
func (a *justLogSwitchAction) Perform(f fritz.HomeAutomationApi) error {
	if a.before.State != a.after.State {
		fmt.Printf("\t'%s'\t%s\t⟶\t%s\n", a.before.Name, console.Btoc(a.before.State), console.Btoc(a.after.State))
	}
	return nil
}

// AhaApiApplier is an Applier that performs changes to the AHA system via the HTTP API.
func AhaApiApplier(f fritz.HomeAutomationApi) Applier {
	return &ahaApiApplier{fritz: f}
}

type ahaApiApplier struct {
	fritz fritz.HomeAutomationApi
}

// Apply does only log the proposed changes.
func (a *ahaApiApplier) Apply(src, target *Plan) error {
	planner := TargetBasedPlanner(reconfigureSwitch, reconfigureThermostat)
	actions, err := planner.Plan(src, target)
	if err != nil {
		return err
	}
	for _, action := range actions {
		action.Perform(a.fritz)
	}
	return nil
}

type reconfigureSwitchAction struct {
	before Switch
	after  Switch
}

func reconfigureSwitch(before, after Switch) Action {
	return &reconfigureSwitchAction{before: before, after: after}
}

// Perform applies the target state to a switch by turning it on/off.
func (a *reconfigureSwitchAction) Perform(f fritz.HomeAutomationApi) error {
	if a.before.State != a.after.State {
		if a.after.State {
			f.SwitchOn(a.before.ain)
		} else {
			f.SwitchOff(a.before.ain)
		}
	}
	return nil
}

type reconfigureThermostatAction struct {
	before Thermostat
	after  Thermostat
}

func reconfigureThermostat(before, after Thermostat) Action {
	return &reconfigureThermostatAction{before: before, after: after}
}

// Perform applies the target state to a switch by turning it on/off.
func (a *reconfigureThermostatAction) Perform(f fritz.HomeAutomationApi) error {
	if a.before.Temperature != a.after.Temperature {
		f.ApplyTemperature(a.after.Temperature, a.before.ain)
	}
	return nil
}
<commit_msg>add logs and return errors<commit_after>package manifest

import (
	"fmt"

	"github.com/bpicode/fritzctl/console"
	"github.com/bpicode/fritzctl/fritz"
)

// Applier defines the interface to apply a plan to the AHA system.
type Applier interface {
	// Apply performs the changes necessary to transition from src to target configuration. If the target plan
	// could be realized, it returns an error. If the plan was applied successfully, it returns null.
	Apply(src, target *Plan) error
}

// DryRunner is an Applier that only plans changes to the AHA system.
func DryRunner() Applier {
	return &dryRunner{}
}

type dryRunner struct {
}

// Apply does only log the proposed changes.
func (d *dryRunner) Apply(src, target *Plan) error {
	planner := TargetBasedPlanner(justLogSwitchState, justLogThermostat)
	actions, err := planner.Plan(src, target)
	if err != nil {
		return err
	}
	fmt.Println("\n\nThe following actions would be applied by the manifest:\n")
	for _, action := range actions {
		action.Perform(nil)
	}
	return nil
}

func justLogThermostat(before, after Thermostat) Action {
	return &justLogThermostatAction{before: before, after: after}
}

type justLogThermostatAction struct {
	before Thermostat
	after  Thermostat
}

// Perform only logs changes.
func (a *justLogThermostatAction) Perform(f fritz.HomeAutomationApi) error {
	if a.before.Temperature != a.after.Temperature {
		fmt.Printf("\t'%s'\t%.1f°C\t⟶\t%.1f°C\n", a.before.Name, a.before.Temperature, a.after.Temperature)
	}
	return nil
}

type justLogSwitchAction struct {
	before Switch
	after  Switch
}

func justLogSwitchState(before, after Switch) Action {
	return &justLogSwitchAction{before: before, after: after}
}

// Perform only logs changes.
func (a *justLogSwitchAction) Perform(f fritz.HomeAutomationApi) error {
	if a.before.State != a.after.State {
		fmt.Printf("\t'%s'\t%s\t⟶\t%s\n", a.before.Name, console.Btoc(a.before.State), console.Btoc(a.after.State))
	}
	return nil
}

// AhaApiApplier is an Applier that performs changes to the AHA system via the HTTP API.
func AhaApiApplier(f fritz.HomeAutomationApi) Applier {
	return &ahaApiApplier{fritz: f}
}

type ahaApiApplier struct {
	fritz fritz.HomeAutomationApi
}

// Apply does only log the proposed changes.
func (a *ahaApiApplier) Apply(src, target *Plan) error {
	planner := TargetBasedPlanner(reconfigureSwitch, reconfigureThermostat)
	actions, err := planner.Plan(src, target)
	if err != nil {
		return err
	}
	for _, action := range actions {
		action.Perform(a.fritz)
	}
	return nil
}

type reconfigureSwitchAction struct {
	before Switch
	after  Switch
}

func reconfigureSwitch(before, after Switch) Action {
	return &reconfigureSwitchAction{before: before, after: after}
}

// Perform applies the target state to a switch by turning it on/off.
func (a *reconfigureSwitchAction) Perform(f fritz.HomeAutomationApi) (err error) {
	if a.before.State != a.after.State {
		if a.after.State {
			_, err = f.SwitchOn(a.before.ain)
		} else {
			_, err = f.SwitchOff(a.before.ain)
		}
		if err == nil {
			fmt.Printf("\tOK\t'%s'\t%s\t⟶\t%s\n", a.before.Name, console.Btoc(a.before.State), console.Btoc(a.after.State))
		}
	}
	return err
}

type reconfigureThermostatAction struct {
	before Thermostat
	after  Thermostat
}

func reconfigureThermostat(before, after Thermostat) Action {
	return &reconfigureThermostatAction{before: before, after: after}
}

// Perform applies the target state to a switch by turning it on/off.
func (a *reconfigureThermostatAction) Perform(f fritz.HomeAutomationApi) (err error) {
	if a.before.Temperature != a.after.Temperature {
		_, err = f.ApplyTemperature(a.after.Temperature, a.before.ain)
		if err == nil {
			fmt.Printf("\tOK\t'%s'\t%.1f°C\t⟶\t%.1f°C\n", a.before.Name, a.before.Temperature, a.after.Temperature)
		}
		return err
	}
	return err
}
<|endoftext|>


<commit_before>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package index

import (
	"reflect"
	"strings"
	"testing"
)

func stringToPath(s string) Path {
	return stringSliceToPath(strings.Split(s, PathSeparator))
}

func stringSliceToPath(s []string) Path {
	p := make(Path, len(s))
	for i, x := range s {
		p[i] = Key(x)
	}
	return p
}

func TestPathEqual(t *testing.T) {
	tests := []struct {
		p, q  Path
		equal bool
	}{
		{
			Path(nil), Path(nil),
			true,
		},
		{
			Path([]Key{}), Path([]Key{}),
			true,
		},
		{
			stringToPath("a"), stringToPath("a"),
			true,
		},
		{
			stringToPath("a:b"), stringToPath("a"),
			false,
		},
		{
			stringToPath("a:b"), stringToPath("a:c"),
			false,
		},
		{
			Path(nil), stringToPath("a"),
			false,
		},
	}

	for ii, tt := range tests {
		if tt.p.Equal(tt.q) != tt.equal {
			t.Errorf("[%d] (%#v).Equal(%#v) = %v, expected %v", ii, tt.p, tt.q, !tt.equal, tt.equal)
		}
	}
}

func TestPathContains(t *testing.T) {
	tests := []struct {
		p, q     Path
		contains bool
	}{
		{
			Path(nil), Path(nil),
			false,
		},
		{
			Path([]Key{}), Path([]Key{}),
			false,
		},
		{
			NewPath("a"), NewPath("a"),
			true,
		},
		{
			NewPath("a:b"), NewPath("a"),
			false,
		},
		{
			NewPath("a:b"), NewPath("a:c"),
			false,
		},
		{
			Path(nil), NewPath("a"),
			false,
		},
		{
			NewPath("a:b"), NewPath("a:b:c"),
			true,
		},
	}

	for ii, tt := range tests {
		if tt.p.Contains(tt.q) != tt.contains {
			t.Errorf("[%d] (%#v).Contains(%#v) = %v, expected %v", ii, tt.p, tt.q, !tt.contains, tt.contains)
		}
	}
}

func TestOrderedIntersection(t *testing.T) {
	tests := []struct {
		in  [][]Path
		out []Path
	}{
		{
			in:  nil,
			out: []Path{},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
			},
			out: []Path{stringToPath("A")},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B")},
			},
			out: []Path{},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B"), stringToPath("A")},
			},
			out: []Path{stringToPath("A")},
		},

		{
			in: [][]Path{
				{stringToPath("A"), stringToPath("B")},
				{stringToPath("B"), stringToPath("A")},
				{stringToPath("A"), stringToPath("B"), stringToPath("C")},
				{stringToPath("C"), stringToPath("A"), stringToPath("B"), stringToPath("B")},
			},
			out: []Path{stringToPath("B"), stringToPath("A")},
		},
	}

	for ii, tt := range tests {
		got := OrderedIntersection(tt.in...)
		if !reflect.DeepEqual(got, tt.out) {
			t.Errorf("[%d] got %#v, expected: %#v", ii, got, tt.out)
		}
	}
}

func TestUnion(t *testing.T) {
	tests := []struct {
		in  [][]Path
		out []Path
	}{
		{
			in:  nil,
			out: []Path{},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
			},
			out: []Path{stringToPath("A")},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B")},
			},
			out: []Path{stringToPath("A"), stringToPath("B")},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B"), stringToPath("A")},
			},
			out: []Path{stringToPath("A"), stringToPath("B")},
		},

		{
			in: [][]Path{
				{stringToPath("A"), stringToPath("B")},
				{stringToPath("B"), stringToPath("A")},
				{stringToPath("A"), stringToPath("B"), stringToPath("C")},
				{stringToPath("C"), stringToPath("A"), stringToPath("B"), stringToPath("B")},
			},
			out: []Path{stringToPath("A"), stringToPath("B"), stringToPath("C")},
		},
	}

	for ii, tt := range tests {
		got := Union(tt.in...)
		if !reflect.DeepEqual(got, tt.out) {
			t.Errorf("[%d] got %#v, expected: %#v", ii, got, tt.out)
		}
	}
}

func TestIndexOfPath(t *testing.T) {
	tests := []struct {
		haystack []Path
		needle   Path
		idx      int
	}{
		{
			[]Path{},
			Path{},
			-1,
		},
		{
			[]Path{
				Path{"Root"},
			},
			Path{"Root"},
			0,
		}, {
			[]Path{
				Path{"Root"}, Path{"Root", "One"},
			},
			Path{"Root", "One"},
			1,
		},
	}

	for ii, tt := range tests {
		i := IndexOfPath(tt.haystack, tt.needle)
		if i != tt.idx {
			t.Errorf("[%d] IndexOfPath(%v, %v) = %d, expected %d", ii, tt.haystack, tt.needle, i, tt.idx)
		}
	}
}

func TestPathFromJSONInterface(t *testing.T) {
	tests := []struct {
		in  interface{}
		out Path
	}{
		{
			interface{}(nil),
			Path(nil),
		},

		{
			[]interface{}{nil},
			Path(nil),
		},

		{
			[]interface{}{""},
			Path{""},
		},
	}

	for ii, tt := range tests {
		got, _ := PathFromJSONInterface(tt.in)
		if !reflect.DeepEqual(got, tt.out) {
			t.Errorf("[%d] got: %#v, expected %#v", ii, got, tt.out)
		}
	}
}
<commit_msg>Improve test coverage for PathFromJSONInterface<commit_after>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package index

import (
	"reflect"
	"strings"
	"testing"
)

func stringToPath(s string) Path {
	return stringSliceToPath(strings.Split(s, PathSeparator))
}

func stringSliceToPath(s []string) Path {
	p := make(Path, len(s))
	for i, x := range s {
		p[i] = Key(x)
	}
	return p
}

func TestPathEqual(t *testing.T) {
	tests := []struct {
		p, q  Path
		equal bool
	}{
		{
			Path(nil), Path(nil),
			true,
		},
		{
			Path([]Key{}), Path([]Key{}),
			true,
		},
		{
			stringToPath("a"), stringToPath("a"),
			true,
		},
		{
			stringToPath("a:b"), stringToPath("a"),
			false,
		},
		{
			stringToPath("a:b"), stringToPath("a:c"),
			false,
		},
		{
			Path(nil), stringToPath("a"),
			false,
		},
	}

	for ii, tt := range tests {
		if tt.p.Equal(tt.q) != tt.equal {
			t.Errorf("[%d] (%#v).Equal(%#v) = %v, expected %v", ii, tt.p, tt.q, !tt.equal, tt.equal)
		}
	}
}

func TestPathContains(t *testing.T) {
	tests := []struct {
		p, q     Path
		contains bool
	}{
		{
			Path(nil), Path(nil),
			false,
		},
		{
			Path([]Key{}), Path([]Key{}),
			false,
		},
		{
			NewPath("a"), NewPath("a"),
			true,
		},
		{
			NewPath("a:b"), NewPath("a"),
			false,
		},
		{
			NewPath("a:b"), NewPath("a:c"),
			false,
		},
		{
			Path(nil), NewPath("a"),
			false,
		},
		{
			NewPath("a:b"), NewPath("a:b:c"),
			true,
		},
	}

	for ii, tt := range tests {
		if tt.p.Contains(tt.q) != tt.contains {
			t.Errorf("[%d] (%#v).Contains(%#v) = %v, expected %v", ii, tt.p, tt.q, !tt.contains, tt.contains)
		}
	}
}

func TestOrderedIntersection(t *testing.T) {
	tests := []struct {
		in  [][]Path
		out []Path
	}{
		{
			in:  nil,
			out: []Path{},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
			},
			out: []Path{stringToPath("A")},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B")},
			},
			out: []Path{},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B"), stringToPath("A")},
			},
			out: []Path{stringToPath("A")},
		},

		{
			in: [][]Path{
				{stringToPath("A"), stringToPath("B")},
				{stringToPath("B"), stringToPath("A")},
				{stringToPath("A"), stringToPath("B"), stringToPath("C")},
				{stringToPath("C"), stringToPath("A"), stringToPath("B"), stringToPath("B")},
			},
			out: []Path{stringToPath("B"), stringToPath("A")},
		},
	}

	for ii, tt := range tests {
		got := OrderedIntersection(tt.in...)
		if !reflect.DeepEqual(got, tt.out) {
			t.Errorf("[%d] got %#v, expected: %#v", ii, got, tt.out)
		}
	}
}

func TestUnion(t *testing.T) {
	tests := []struct {
		in  [][]Path
		out []Path
	}{
		{
			in:  nil,
			out: []Path{},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
			},
			out: []Path{stringToPath("A")},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B")},
			},
			out: []Path{stringToPath("A"), stringToPath("B")},
		},

		{
			in: [][]Path{
				{stringToPath("A")},
				{stringToPath("B"), stringToPath("A")},
			},
			out: []Path{stringToPath("A"), stringToPath("B")},
		},

		{
			in: [][]Path{
				{stringToPath("A"), stringToPath("B")},
				{stringToPath("B"), stringToPath("A")},
				{stringToPath("A"), stringToPath("B"), stringToPath("C")},
				{stringToPath("C"), stringToPath("A"), stringToPath("B"), stringToPath("B")},
			},
			out: []Path{stringToPath("A"), stringToPath("B"), stringToPath("C")},
		},
	}

	for ii, tt := range tests {
		got := Union(tt.in...)
		if !reflect.DeepEqual(got, tt.out) {
			t.Errorf("[%d] got %#v, expected: %#v", ii, got, tt.out)
		}
	}
}

func TestIndexOfPath(t *testing.T) {
	tests := []struct {
		haystack []Path
		needle   Path
		idx      int
	}{
		{
			[]Path{},
			Path{},
			-1,
		},
		{
			[]Path{
				Path{"Root"},
			},
			Path{"Root"},
			0,
		}, {
			[]Path{
				Path{"Root"}, Path{"Root", "One"},
			},
			Path{"Root", "One"},
			1,
		},
	}

	for ii, tt := range tests {
		i := IndexOfPath(tt.haystack, tt.needle)
		if i != tt.idx {
			t.Errorf("[%d] IndexOfPath(%v, %v) = %d, expected %d", ii, tt.haystack, tt.needle, i, tt.idx)
		}
	}
}

func TestPathFromJSONInterface(t *testing.T) {
	tests := []struct {
		in  interface{}
		out Path
	}{
		{
			interface{}(nil),
			Path(nil),
		},

		{
			[]interface{}{nil},
			Path(nil),
		},

		{
			[]interface{}{""},
			Path{""},
		},

		{
			[]interface{}{"", 123.},
			Path{"", "123"},
		},
	}

	for ii, tt := range tests {
		got, _ := PathFromJSONInterface(tt.in)
		if !reflect.DeepEqual(got, tt.out) {
			t.Errorf("[%d] got: %#v, expected %#v", ii, got, tt.out)
		}
	}
}
<|endoftext|>


<commit_before>package bench

import (
	"fmt"

	"flag"
	"io/ioutil"

	"context"
	"math"
	"math/rand"
	"time"
)

// ZipfSetBits sets random bits according to the Zipf-Mandelbrot distribution.
// This distribution accepts two parameters, Exponent and Ratio, for both bitmaps and profiles.
// It also uses PermutationGenerator to permute IDs randomly.
type ZipfSetBits struct {
	HasClient
	Name           string                `json:"name"`
	BaseBitmapID   int64                 `json:"base-bitmap-id"`
	BaseProfileID  int64                 `json:"base-profile-id"`
	BitmapIDRange  int64                 `json:"bitmap-id-range"`
	ProfileIDRange int64                 `json:"profile-id-range"`
	Iterations     int                   `json:"iterations"`
	Seed           int64                 `json:"seed"`
	DB             string                `json:"db"`
	bitmapRng      *rand.Zipf            `json:"-"`
	ProfileRng     *rand.Zipf            `json:"-"`
	BitmapPerm     *PermutationGenerator `json:"-"`
	ProfilePerm    *PermutationGenerator `json:"-"`

	BitmapExponent  float64 `json:"bitmap-exponent"`
	BitmapRatio     float64 `json:"bitmap-ratio"`
	ProfileExponent float64 `json:"profile-exponent"`
	ProfileRatio    float64 `json:"profile-ratio"`
}

func (b *ZipfSetBits) Usage() string {
	return `
zipf-set-bits sets random bits according to the Zipf distribution.
This is a power-law distribution controlled by two parameters.
Exponent, in the range (1, inf), with a default value of 1.001, controls
the "sharpness" of the distribution, with higher exponent being sharper.
Ratio, in the range (0, 1), with a default value of 0.25, controls the
maximum variation of the distribution, with higher ratio being more uniform.

Usage: zipf-set-bits [arguments]

The following arguments are available:

	-base-bitmap-id int
		bits being set will all be greater than BaseBitmapID

	-bitmap-id-range int
		number of possible bitmap ids that can be set

	-base-profile-id int
		profile id num to start from

	-profile-id-range int
		number of possible profile ids that can be set

	-iterations int
		number of bits to set

	-seed int
		Seed for RNG

	-db string
		pilosa db to use

	-bitmap-exponent float64
		zipf exponent parameter for bitmap IDs

	-bitmap-ratio float64
		zipf probability ratio parameter for bitmap IDs

	-profile-exponent float64
		zipf exponent parameter for profile IDs

	-profile-ratio float64
		zipf probability ratio parameter for profile IDs

	-client-type string
		Can be 'single' (all agents hitting one host) or 'round_robin'
`[1:]
}

func (b *ZipfSetBits) ConsumeFlags(args []string) ([]string, error) {
	fs := flag.NewFlagSet("ZipfSetBits", flag.ContinueOnError)
	fs.SetOutput(ioutil.Discard)
	fs.Int64Var(&b.BaseBitmapID, "base-bitmap-id", 0, "")
	fs.Int64Var(&b.BitmapIDRange, "bitmap-id-range", 100000, "")
	fs.Int64Var(&b.BaseProfileID, "base-profile-id", 0, "")
	fs.Int64Var(&b.ProfileIDRange, "profile-id-range", 100000, "")
	fs.Int64Var(&b.Seed, "seed", 1, "")
	fs.IntVar(&b.Iterations, "iterations", 100, "")
	fs.StringVar(&b.DB, "db", "benchdb", "")
	fs.Float64Var(&b.BitmapExponent, "bitmap-exponent", 1.01, "")
	fs.Float64Var(&b.BitmapRatio, "bitmap-ratio", 0.25, "")
	fs.Float64Var(&b.ProfileExponent, "profile-exponent", 1.01, "")
	fs.Float64Var(&b.ProfileRatio, "profile-ratio", 0.25, "")
	fs.StringVar(&b.ClientType, "client-type", "single", "")

	if err := fs.Parse(args); err != nil {
		return nil, err
	}
	return fs.Args(), nil
}

// Offset is the true parameter used by the Zipf distribution, but the ratio,
// as defined here, is a simpler, readable way to define the distribution.
// Offset is in [1, inf), and its meaning depends on N (a pain for updating benchmark configs)
// ratio is in (0, 1), and its meaning does not depend on N.
// it is the ratio of the lowest probability in the distribution to the highest.
// ratio=0.01 corresponds to a very small offset - the most skewed distribution for a given pair (N, exp)
// ratio=0.99 corresponds to a very large offset - the most nearly uniform distribution for a given (N, exp)
func getZipfOffset(N int64, exp, ratio float64) float64 {
	z := math.Pow(ratio, 1/exp)
	return z * float64(N-1) / (1 - z)
}

func (b *ZipfSetBits) Init(hosts []string, agentNum int) error {
	b.Name = "ZipfSetBits"
	rnd := rand.New(rand.NewSource(b.Seed + int64(agentNum)))
	bitmapOffset := getZipfOffset(b.BitmapIDRange, b.BitmapExponent, b.BitmapRatio)
	b.bitmapRng = rand.NewZipf(rnd, b.BitmapExponent, bitmapOffset, uint64(b.BitmapIDRange-1))
	profileOffset := getZipfOffset(b.ProfileIDRange, b.ProfileExponent, b.ProfileRatio)
	b.profileRng = rand.NewZipf(rnd, b.ProfileExponent, profileOffset, uint64(b.ProfileIDRange-1))

	b.bitmapPerm = NewPermutationGenerator(b.BitmapIDRange, b.Seed)
	b.profilePerm = NewPermutationGenerator(b.ProfileIDRange, b.Seed+1)

	return b.HasClient.Init(hosts, agentNum)
}

// Run runs the ZipfSetBits benchmark
func (b *ZipfSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
	results := make(map[string]interface{})
	if b.cli == nil {
		results["error"] = fmt.Errorf("No client set for ZipfSetBits agent: %v", agentNum)
		return results
	}
	s := NewStats()
	var start time.Time
	for n := 0; n < b.Iterations; n++ {
		// generate IDs from Zipf distribution
		bitmapIDOriginal := b.bitmapRng.Uint64()
		profIDOriginal := b.profileRng.Uint64()
		// permute IDs randomly, but repeatably
		bitmapID := b.bitmapPerm.Next(int64(bitmapIDOriginal))
		profID := b.profilePerm.Next(int64(profIDOriginal))

		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+int64(bitmapID), b.BaseProfileID+int64(profID))
		start = time.Now()
		b.cli.ExecuteQuery(ctx, b.DB, query, true)
		s.Add(time.Now().Sub(start))
	}
	AddToResults(s, results)
	return results
}
<commit_msg>Remove extraneous JSON tags<commit_after>package bench

import (
	"fmt"

	"flag"
	"io/ioutil"

	"context"
	"math"
	"math/rand"
	"time"
)

// ZipfSetBits sets random bits according to the Zipf-Mandelbrot distribution.
// This distribution accepts two parameters, Exponent and Ratio, for both bitmaps and profiles.
// It also uses PermutationGenerator to permute IDs randomly.
type ZipfSetBits struct {
	HasClient
	Name            string  `json:"name"`
	BaseBitmapID    int64   `json:"base-bitmap-id"`
	BaseProfileID   int64   `json:"base-profile-id"`
	BitmapIDRange   int64   `json:"bitmap-id-range"`
	ProfileIDRange  int64   `json:"profile-id-range"`
	Iterations      int     `json:"iterations"`
	Seed            int64   `json:"seed"`
	DB              string  `json:"db"`
	BitmapExponent  float64 `json:"bitmap-exponent"`
	BitmapRatio     float64 `json:"bitmap-ratio"`
	ProfileExponent float64 `json:"profile-exponent"`
	ProfileRatio    float64 `json:"profile-ratio"`
	bitmapRng       *rand.Zipf
	ProfileRng      *rand.Zipf
	BitmapPerm      *PermutationGenerator
	ProfilePerm     *PermutationGenerator
}

func (b *ZipfSetBits) Usage() string {
	return `
zipf-set-bits sets random bits according to the Zipf distribution.
This is a power-law distribution controlled by two parameters.
Exponent, in the range (1, inf), with a default value of 1.001, controls
the "sharpness" of the distribution, with higher exponent being sharper.
Ratio, in the range (0, 1), with a default value of 0.25, controls the
maximum variation of the distribution, with higher ratio being more uniform.

Usage: zipf-set-bits [arguments]

The following arguments are available:

	-base-bitmap-id int
		bits being set will all be greater than BaseBitmapID

	-bitmap-id-range int
		number of possible bitmap ids that can be set

	-base-profile-id int
		profile id num to start from

	-profile-id-range int
		number of possible profile ids that can be set

	-iterations int
		number of bits to set

	-seed int
		Seed for RNG

	-db string
		pilosa db to use

	-bitmap-exponent float64
		zipf exponent parameter for bitmap IDs

	-bitmap-ratio float64
		zipf probability ratio parameter for bitmap IDs

	-profile-exponent float64
		zipf exponent parameter for profile IDs

	-profile-ratio float64
		zipf probability ratio parameter for profile IDs

	-client-type string
		Can be 'single' (all agents hitting one host) or 'round_robin'
`[1:]
}

func (b *ZipfSetBits) ConsumeFlags(args []string) ([]string, error) {
	fs := flag.NewFlagSet("ZipfSetBits", flag.ContinueOnError)
	fs.SetOutput(ioutil.Discard)
	fs.Int64Var(&b.BaseBitmapID, "base-bitmap-id", 0, "")
	fs.Int64Var(&b.BitmapIDRange, "bitmap-id-range", 100000, "")
	fs.Int64Var(&b.BaseProfileID, "base-profile-id", 0, "")
	fs.Int64Var(&b.ProfileIDRange, "profile-id-range", 100000, "")
	fs.Int64Var(&b.Seed, "seed", 1, "")
	fs.IntVar(&b.Iterations, "iterations", 100, "")
	fs.StringVar(&b.DB, "db", "benchdb", "")
	fs.Float64Var(&b.BitmapExponent, "bitmap-exponent", 1.01, "")
	fs.Float64Var(&b.BitmapRatio, "bitmap-ratio", 0.25, "")
	fs.Float64Var(&b.ProfileExponent, "profile-exponent", 1.01, "")
	fs.Float64Var(&b.ProfileRatio, "profile-ratio", 0.25, "")
	fs.StringVar(&b.ClientType, "client-type", "single", "")

	if err := fs.Parse(args); err != nil {
		return nil, err
	}
	return fs.Args(), nil
}

// Offset is the true parameter used by the Zipf distribution, but the ratio,
// as defined here, is a simpler, readable way to define the distribution.
// Offset is in [1, inf), and its meaning depends on N (a pain for updating benchmark configs)
// ratio is in (0, 1), and its meaning does not depend on N.
// it is the ratio of the lowest probability in the distribution to the highest.
// ratio=0.01 corresponds to a very small offset - the most skewed distribution for a given pair (N, exp)
// ratio=0.99 corresponds to a very large offset - the most nearly uniform distribution for a given (N, exp)
func getZipfOffset(N int64, exp, ratio float64) float64 {
	z := math.Pow(ratio, 1/exp)
	return z * float64(N-1) / (1 - z)
}

func (b *ZipfSetBits) Init(hosts []string, agentNum int) error {
	b.Name = "ZipfSetBits"
	rnd := rand.New(rand.NewSource(b.Seed + int64(agentNum)))
	bitmapOffset := getZipfOffset(b.BitmapIDRange, b.BitmapExponent, b.BitmapRatio)
	b.bitmapRng = rand.NewZipf(rnd, b.BitmapExponent, bitmapOffset, uint64(b.BitmapIDRange-1))
	profileOffset := getZipfOffset(b.ProfileIDRange, b.ProfileExponent, b.ProfileRatio)
	b.profileRng = rand.NewZipf(rnd, b.ProfileExponent, profileOffset, uint64(b.ProfileIDRange-1))

	b.bitmapPerm = NewPermutationGenerator(b.BitmapIDRange, b.Seed)
	b.profilePerm = NewPermutationGenerator(b.ProfileIDRange, b.Seed+1)

	return b.HasClient.Init(hosts, agentNum)
}

// Run runs the ZipfSetBits benchmark
func (b *ZipfSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
	results := make(map[string]interface{})
	if b.cli == nil {
		results["error"] = fmt.Errorf("No client set for ZipfSetBits agent: %v", agentNum)
		return results
	}
	s := NewStats()
	var start time.Time
	for n := 0; n < b.Iterations; n++ {
		// generate IDs from Zipf distribution
		bitmapIDOriginal := b.bitmapRng.Uint64()
		profIDOriginal := b.profileRng.Uint64()
		// permute IDs randomly, but repeatably
		bitmapID := b.bitmapPerm.Next(int64(bitmapIDOriginal))
		profID := b.profilePerm.Next(int64(profIDOriginal))

		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+int64(bitmapID), b.BaseProfileID+int64(profID))
		start = time.Now()
		b.cli.ExecuteQuery(ctx, b.DB, query, true)
		s.Add(time.Now().Sub(start))
	}
	AddToResults(s, results)
	return results
}
<|endoftext|>


<commit_before>package metrics

import (
	"fmt"
	"net/http"
	"runtime"
	"sort"
)

var prefix = ""

func SetPrefix(p string) {
	prefix = p
}

var memstats = new(runtime.MemStats)

func init() {
	http.Handle("/metrics", http.HandlerFunc(printMetrics))
}

func printMetrics(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")

	// Runtime memory stats
	runtime.ReadMemStats(memstats)
	// General statistics.
	fmt.Fprintf(w, "%smem_alloc %d\n", prefix, memstats.Alloc)            // bytes allocated and not yet freed
	fmt.Fprintf(w, "%smem_alloc_total %d\n", prefix, memstats.TotalAlloc) // bytes allocated (even if freed)
	fmt.Fprintf(w, "%smem_sys %d\n", prefix, memstats.Sys)                // bytes obtained from system (sum of XxxSys below)
	fmt.Fprintf(w, "%smem_ptr_lookups %d\n", prefix, memstats.Lookups)    // number of pointer lookups
	fmt.Fprintf(w, "%smem_mallocs %d\n", prefix, memstats.Mallocs)        // number of mallocs
	fmt.Fprintf(w, "%smem_frees %d\n", prefix, memstats.Frees)            // number of frees

	// Main allocation heap statistics.
	fmt.Fprintf(w, "%smem_heap_alloc %d\n", prefix, memstats.HeapAlloc)       // bytes allocated and not yet freed (same as Alloc above)
	fmt.Fprintf(w, "%smem_heap_sys %d\n", prefix, memstats.HeapSys)           // bytes obtained from system
	fmt.Fprintf(w, "%smem_heap_idle %d\n", prefix, memstats.HeapIdle)         // bytes in idle spans
	fmt.Fprintf(w, "%smem_heap_in_use %d\n", prefix, memstats.HeapInuse)      // bytes in non-idle span
	fmt.Fprintf(w, "%smem_heap_released %d\n", prefix, memstats.HeapReleased) // bytes released to the OS
	fmt.Fprintf(w, "%smem_heap_objects %d\n", prefix, memstats.HeapObjects)   // total number of allocated objects

	fmt.Fprintf(w, "%smem_stack_in_use %d\n", prefix, memstats.StackInuse) // bytes used by stack allocator
	fmt.Fprintf(w, "%smem_stack_sys %d\n", prefix, memstats.StackSys)
	fmt.Fprintf(w, "%smem_mspan_in_use %d\n", prefix, memstats.MSpanInuse) // mspan structures
	fmt.Fprintf(w, "%smem_mspan_sys %d\n", prefix, memstats.MSpanSys)
	fmt.Fprintf(w, "%smem_mcache_in_use %d\n", prefix, memstats.MCacheInuse) // mcache structures
	fmt.Fprintf(w, "%smem_mcache_sys %d\n", prefix, memstats.MCacheSys)
	fmt.Fprintf(w, "%smem_buck_hash_sys %d\n", prefix, memstats.BuckHashSys) // profiling bucket hash table
	fmt.Fprintf(w, "%smem_gc_sys %d\n", prefix, memstats.GCSys)              // GC metadata
	fmt.Fprintf(w, "%smem_other_sys %d\n", prefix, memstats.OtherSys)        // other system allocations

	fmt.Fprintf(w, "%sgc_next_gc_heap_alloc %d\n", prefix, memstats.NextGC) // next collection will happen when HeapAlloc ≥ this amount
	fmt.Fprintf(w, "%sgc_last_gc_time %d\n", prefix, memstats.LastGC)       // end time of last collection (nanoseconds since 1970)
	fmt.Fprintf(w, "%sgc_pause_total %d\n", prefix, memstats.PauseTotalNs)
	fmt.Fprintf(w, "%sgc_num_gc %d\n", prefix, memstats.NumGC)
	fmt.Fprintf(w, "%sgc_gc_cpu_frac %f\n", prefix, memstats.GCCPUFraction)

	// circular buffer of recent GC pause durations, most recent at [(NumGC+255)%256]
	for i, p := range pausePercentiles(memstats.PauseNs[:], memstats.NumGC) {
		fmt.Fprintf(w, "%sgc_pause_pctl_%d %d\n", prefix, i*5, p)
	}

	// Per-size allocation statistics.
	for _, b := range memstats.BySize {
		fmt.Fprintf(w, "%salloc_size_%d_mallocs %d\n", prefix, b.Size, b.Mallocs)
		fmt.Fprintf(w, "%salloc_size_%d_frees %d\n", prefix, b.Size, b.Frees)
	}

	// Histograms
	hists := getAllHistograms()
	for name, dat := range hists {
		fmt.Fprintf(w, "%shist_%s_count %d\n", prefix, name, *dat.count)
		fmt.Fprintf(w, "%shist_%s_kept %d\n", prefix, name, *dat.kept)
		for i, p := range hdatPercentiles(dat) {
			fmt.Fprintf(w, "%shist_%s_pctl_%d %d\n", prefix, name, i*5, p)
		}
	}

	// Application stats
	ctrs := getAllCounters()
	for name, val := range ctrs {
		fmt.Fprintf(w, "%s%s %d\n", prefix, name, val)
	}
}

// Percentiles go by 5% percentile steps from min to max. We report all of them even though it's
// likely only min, 25th, 50th, 75th, 95th, and max will be used. It's assumed the metric poller
// that is consuming this output will choose to only report to the metrics system what it considers
// useful information.
//
// Slice layout:
//  [0]: min (0th)
//  [1]: 5th
//  [n]: 5n
//  [19]: 95th
//  [20]: max (100th)
func hdatPercentiles(dat *hdat) []uint64 {
	buf := dat.buf
	kept := *dat.kept

	if kept == 0 {
		return nil
	}
	if kept < uint64(len(buf)) {
		buf = buf[:kept]
	}

	sort.Sort(uint64slice(buf))

	pctls := make([]uint64, 21)
	pctls[0] = *dat.min
	pctls[20] = *dat.max

	for i := 1; i < 20; i++ {
		idx := len(buf) * i / 20
		pctls[i] = buf[idx]
	}

	return pctls
}

func pausePercentiles(pauses []uint64, ngc uint32) []uint64 {
	if ngc < uint32(len(pauses)) {
		pauses = pauses[:ngc]
	}

	sort.Sort(uint64slice(pauses))

	pctls := make([]uint64, 21)
	pctls[0] = pauses[0]
	pctls[20] = pauses[len(pauses)-1]

	for i := 1; i < 20; i++ {
		idx := len(pauses) * i / 20
		pctls[i] = pauses[idx]
	}

	return pctls
}

type uint64slice []uint64

func (u uint64slice) Len() int           { return len(u) }
func (u uint64slice) Swap(i, j int)      { u[i], u[j] = u[j], u[i] }
func (u uint64slice) Less(i, j int) bool { return u[i] < u[j] }
<commit_msg>Adding 99th percentiles for all histograms and GC times<commit_after>package metrics

import (
	"fmt"
	"net/http"
	"runtime"
	"sort"
)

var prefix = ""

func SetPrefix(p string) {
	prefix = p
}

var memstats = new(runtime.MemStats)

func init() {
	http.Handle("/metrics", http.HandlerFunc(printMetrics))
}

func printMetrics(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")

	// Runtime memory stats
	runtime.ReadMemStats(memstats)
	// General statistics.
	fmt.Fprintf(w, "%smem_alloc %d\n", prefix, memstats.Alloc)            // bytes allocated and not yet freed
	fmt.Fprintf(w, "%smem_alloc_total %d\n", prefix, memstats.TotalAlloc) // bytes allocated (even if freed)
	fmt.Fprintf(w, "%smem_sys %d\n", prefix, memstats.Sys)                // bytes obtained from system (sum of XxxSys below)
	fmt.Fprintf(w, "%smem_ptr_lookups %d\n", prefix, memstats.Lookups)    // number of pointer lookups
	fmt.Fprintf(w, "%smem_mallocs %d\n", prefix, memstats.Mallocs)        // number of mallocs
	fmt.Fprintf(w, "%smem_frees %d\n", prefix, memstats.Frees)            // number of frees

	// Main allocation heap statistics.
	fmt.Fprintf(w, "%smem_heap_alloc %d\n", prefix, memstats.HeapAlloc)       // bytes allocated and not yet freed (same as Alloc above)
	fmt.Fprintf(w, "%smem_heap_sys %d\n", prefix, memstats.HeapSys)           // bytes obtained from system
	fmt.Fprintf(w, "%smem_heap_idle %d\n", prefix, memstats.HeapIdle)         // bytes in idle spans
	fmt.Fprintf(w, "%smem_heap_in_use %d\n", prefix, memstats.HeapInuse)      // bytes in non-idle span
	fmt.Fprintf(w, "%smem_heap_released %d\n", prefix, memstats.HeapReleased) // bytes released to the OS
	fmt.Fprintf(w, "%smem_heap_objects %d\n", prefix, memstats.HeapObjects)   // total number of allocated objects

	fmt.Fprintf(w, "%smem_stack_in_use %d\n", prefix, memstats.StackInuse) // bytes used by stack allocator
	fmt.Fprintf(w, "%smem_stack_sys %d\n", prefix, memstats.StackSys)
	fmt.Fprintf(w, "%smem_mspan_in_use %d\n", prefix, memstats.MSpanInuse) // mspan structures
	fmt.Fprintf(w, "%smem_mspan_sys %d\n", prefix, memstats.MSpanSys)
	fmt.Fprintf(w, "%smem_mcache_in_use %d\n", prefix, memstats.MCacheInuse) // mcache structures
	fmt.Fprintf(w, "%smem_mcache_sys %d\n", prefix, memstats.MCacheSys)
	fmt.Fprintf(w, "%smem_buck_hash_sys %d\n", prefix, memstats.BuckHashSys) // profiling bucket hash table
	fmt.Fprintf(w, "%smem_gc_sys %d\n", prefix, memstats.GCSys)              // GC metadata
	fmt.Fprintf(w, "%smem_other_sys %d\n", prefix, memstats.OtherSys)        // other system allocations

	fmt.Fprintf(w, "%sgc_next_gc_heap_alloc %d\n", prefix, memstats.NextGC) // next collection will happen when HeapAlloc ≥ this amount
	fmt.Fprintf(w, "%sgc_last_gc_time %d\n", prefix, memstats.LastGC)       // end time of last collection (nanoseconds since 1970)
	fmt.Fprintf(w, "%sgc_pause_total %d\n", prefix, memstats.PauseTotalNs)
	fmt.Fprintf(w, "%sgc_num_gc %d\n", prefix, memstats.NumGC)
	fmt.Fprintf(w, "%sgc_gc_cpu_frac %f\n", prefix, memstats.GCCPUFraction)

	// circular buffer of recent GC pause durations, most recent at [(NumGC+255)%256]
	pctls := pausePercentiles(memstats.PauseNs[:], memstats.NumGC)
	for i := 0; i < 20; i++ {
		p := pctls[i]
		fmt.Fprintf(w, "%sgc_pause_pctl_%d %d\n", prefix, i*5, p)
	}
	fmt.Fprintf(w, "%sgc_pause_pctl_%d %d\n", prefix, 99, pctls[20])
	fmt.Fprintf(w, "%sgc_pause_pctl_%d %d\n", prefix, 100, pctls[21])

	// Per-size allocation statistics.
	for _, b := range memstats.BySize {
		fmt.Fprintf(w, "%salloc_size_%d_mallocs %d\n", prefix, b.Size, b.Mallocs)
		fmt.Fprintf(w, "%salloc_size_%d_frees %d\n", prefix, b.Size, b.Frees)
	}

	// Histograms
	hists := getAllHistograms()
	for name, dat := range hists {
		fmt.Fprintf(w, "%shist_%s_count %d\n", prefix, name, *dat.count)
		fmt.Fprintf(w, "%shist_%s_kept %d\n", prefix, name, *dat.kept)
		pctls := hdatPercentiles(dat)
		for i := 0; i < 20; i++ {
			p := pctls[i]
			fmt.Fprintf(w, "%shist_%s_pctl_%d %d\n", prefix, name, i*5, p)
		}
		fmt.Fprintf(w, "%shist_%s_pctl_%d %d\n", prefix, name, 99, pctls[20])
		fmt.Fprintf(w, "%shist_%s_pctl_%d %d\n", prefix, name, 100, pctls[21])
	}

	// Application stats
	ctrs := getAllCounters()
	for name, val := range ctrs {
		fmt.Fprintf(w, "%s%s %d\n", prefix, name, val)
	}
}

// Percentiles go by 5% percentile steps from min to max. We report all of them even though it's
// likely only min, 25th, 50th, 75th, 95th, 99th, and max will be used. It's assumed the metric
// poller that is consuming this output will choose to only report to the metrics system what it
// considers useful information.
//
// Slice layout:
//  [0]: min (0th)
//  [1]: 5th
//  [n]: 5n
//  [19]: 95th
//  [20]: 99th
//  [21]: max (100th)
func hdatPercentiles(dat *hdat) []uint64 {
	buf := dat.buf
	kept := *dat.kept

	if kept == 0 {
		return nil
	}
	if kept < uint64(len(buf)) {
		buf = buf[:kept]
	}

	sort.Sort(uint64slice(buf))

	pctls := make([]uint64, 22)
	pctls[0] = *dat.min
	pctls[21] = *dat.max

	for i := 1; i < 20; i++ {
		idx := len(buf) * i / 20
		pctls[i] = buf[idx]
	}

	idx := len(buf) * 99 / 100
	pctls[20] = buf[idx]

	return pctls
}

func pausePercentiles(pauses []uint64, ngc uint32) []uint64 {
	if ngc < uint32(len(pauses)) {
		pauses = pauses[:ngc]
	}

	sort.Sort(uint64slice(pauses))

	pctls := make([]uint64, 22)
	pctls[0] = pauses[0]
	pctls[21] = pauses[len(pauses)-1]

	for i := 1; i < 20; i++ {
		idx := len(pauses) * i / 20
		pctls[i] = pauses[idx]
	}

	idx := len(pauses) * 99 / 100
	pctls[20] = pauses[idx]

	return pctls
}

type uint64slice []uint64

func (u uint64slice) Len() int           { return len(u) }
func (u uint64slice) Swap(i, j int)      { u[i], u[j] = u[j], u[i] }
func (u uint64slice) Less(i, j int) bool { return u[i] < u[j] }
<|endoftext|>


<commit_before>// Copyright 2019 The grok_exporter Authors
//
// 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 fswatcher

import (
	"fmt"
	"github.com/sirupsen/logrus"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"
)

type watcher struct {
	fd int
}

type fileWithReader struct {
	file   *os.File
	reader *lineReader
}

func (w *watcher) unwatchDir(dir *Dir) error {
	// After calling eventProducerLoop.Close(), we need to call inotify_rm_watch()
	// in order to terminate the inotify loop. See eventProducerLoop.Close().
	success, err := syscall.InotifyRmWatch(w.fd, uint32(dir.wd))
	if success != 0 || err != nil {
		return fmt.Errorf("inotify_rm_watch(%q) failed: status=%v, err=%v", dir.path, success, err)
	} else {
		return nil
	}
}

func (w *watcher) Close() error {
	err := syscall.Close(w.fd)
	if err != nil {
		return fmt.Errorf("failed to close the inotify file descriptor: %v", err)
	} else {
		return nil
	}
}

func unwatchDirByEvent(t *fileTailer, event inotifyEvent) {
	watchedDirsAfter := make([]*Dir, 0, len(t.watchedDirs)-1)
	for _, existing := range t.watchedDirs {
		if existing.wd != int(event.Wd) {
			watchedDirsAfter = append(watchedDirsAfter, existing)
		}
	}
	t.watchedDirs = watchedDirsAfter
}

func (w *watcher) runFseventProducerLoop() fseventProducerLoop {
	return runInotifyLoop(w.fd)
}

func initWatcher() (fswatcher, Error) {
	fd, err := syscall.InotifyInit1(syscall.IN_CLOEXEC)
	if err != nil {
		return nil, NewError(NotSpecified, err, "inotify_init1() failed")
	}
	return &watcher{fd: fd}, nil
}

func (w *watcher) watchDir(path string) (*Dir, Error) {
	var (
		dir *Dir
		err error
		Err Error
	)
	dir, Err = newDir(path)
	if Err != nil {
		return nil, Err
	}
	dir.wd, err = syscall.InotifyAddWatch(w.fd, path, syscall.IN_MODIFY|syscall.IN_MOVED_FROM|syscall.IN_MOVED_TO|syscall.IN_DELETE|syscall.IN_CREATE)
	if err != nil {
		return nil, NewErrorf(NotSpecified, err, "%q: inotify_add_watch() failed", path)
	}
	return dir, nil
}

func newDir(path string) (*Dir, Error) {
	return &Dir{path: path}, nil
}

func (w *watcher) watchFile(_ fileMeta) Error {
	// nothing to do, because on Linux we watch the directory and don't need to watch individual files.
	return nil
}

func findDir(t *fileTailer, event inotifyEvent) (*Dir, Error) {
	for _, dir := range t.watchedDirs {
		if dir.wd == int(event.Wd) {
			return dir, nil
		}
	}
	// We should always find the directory in watchedDirs if we receive an inotifyEvent.
	// The following code should never be executed.
	// However, Github issue #124 is a bug report where we receive an inotifyEvent for an unwatched directory.
	// Create a detailed error message to help figuring out why this happens.
	var (
		errMsg strings.Builder
		first  = true
	)
	fmt.Fprintf(&errMsg, "watch list inconsistent: cannot find directory for file system event %q in the list of watched directories [", event.String())
	for _, dir := range t.watchedDirs {
		if !first {
			fmt.Fprintf(&errMsg, ", ")
		}
		fmt.Fprintf(&errMsg, "%q", dir.Path())
		first = false
	}
	fmt.Fprintf(&errMsg, "]")
	return nil, NewError(NotSpecified, nil, errMsg.String())
}

func (w *watcher) processEvent(t *fileTailer, fsevent fsevent, log logrus.FieldLogger) Error {
	event, ok := fsevent.(inotifyEvent)
	if !ok {
		return NewErrorf(NotSpecified, nil, "received a file system event of unknown type %T", event)
	}
	dir, Err := findDir(t, event)
	if Err != nil {
		return Err
	}
	dirLogger := log.WithField("directory", dir.path)
	dirLogger.Debugf("received event: %v", event)
	if event.Mask&syscall.IN_IGNORED == syscall.IN_IGNORED {
		unwatchDirByEvent(t, event) // need to remove it from watchedDirs, because otherwise we close the removed dir on shutdown which causes an error
		return NewErrorf(NotSpecified, nil, "%s: directory was removed while being watched", dir.path)
	}
	if event.Mask&syscall.IN_MODIFY == syscall.IN_MODIFY {
		file, ok := t.watchedFiles[filepath.Join(dir.path, event.Name)]
		if !ok {
			return nil // unrelated file was modified
		}
		truncated, err := isTruncated(file.file)
		if err != nil {
			return NewErrorf(NotSpecified, err, "%v: seek() or stat() failed", file.file.Name())
		}
		if truncated {
			_, err = file.file.Seek(0, io.SeekStart)
			if err != nil {
				return NewErrorf(NotSpecified, err, "%v: seek() failed", file.file.Name())
			}
			file.reader.Clear()
		}
		readErr := t.readNewLines(file, dirLogger)
		if readErr != nil {
			return readErr
		}
	}
	if event.Mask&syscall.IN_MOVED_FROM == syscall.IN_MOVED_FROM || event.Mask&syscall.IN_DELETE == syscall.IN_DELETE || event.Mask&syscall.IN_CREATE == syscall.IN_CREATE || event.Mask&syscall.IN_MOVED_TO == syscall.IN_MOVED_TO {
		// There are a lot of corner cases here:
		// * a file is renamed, but still matches the pattern so we continue watching it (MOVED_FROM followed by MOVED_TO)
		// * a file is created overwriting an existing file
		// * a file is moved to the watched directory overwriting an existing file
		// Trying to figure out what happened from the events would be error prone.
		// Therefore, we don't care which of the above events we received, we just update our watched files with the current
		// state of the watched directory.
		err := t.syncFilesInDir(dir, true, dirLogger)
		if err != nil {
			return err
		}
	}
	return nil
}

func isTruncated(file *os.File) (bool, error) {
	currentPos, err := file.Seek(0, io.SeekCurrent)
	if err != nil {
		return false, err
	}
	fileInfo, err := file.Stat()
	if err != nil {
		return false, err
	}
	return currentPos > fileInfo.Size(), nil
}

func findSameFile(t *fileTailer, file os.FileInfo, _ string) (*fileWithReader, Error) {
	var (
		fileInfo os.FileInfo
		err      error
	)
	for _, watchedFile := range t.watchedFiles {
		fileInfo, err = watchedFile.file.Stat()
		if err != nil {
			return nil, NewErrorf(NotSpecified, err, "%v: stat failed", watchedFile.file.Name())
		}
		if os.SameFile(fileInfo, file) {
			return watchedFile, nil
		}
	}
	return nil, nil
}
<commit_msg>#124 improved error message for file system watcher<commit_after>// Copyright 2019 The grok_exporter Authors
//
// 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 fswatcher

import (
	"fmt"
	"github.com/sirupsen/logrus"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"
)

type watcher struct {
	fd int
}

type fileWithReader struct {
	file   *os.File
	reader *lineReader
}

func (w *watcher) unwatchDir(dir *Dir) error {
	// After calling eventProducerLoop.Close(), we need to call inotify_rm_watch()
	// in order to terminate the inotify loop. See eventProducerLoop.Close().
	success, err := syscall.InotifyRmWatch(w.fd, uint32(dir.wd))
	if success != 0 || err != nil {
		return fmt.Errorf("inotify_rm_watch(%q) failed: status=%v, err=%v", dir.path, success, err)
	} else {
		return nil
	}
}

func (w *watcher) Close() error {
	err := syscall.Close(w.fd)
	if err != nil {
		return fmt.Errorf("failed to close the inotify file descriptor: %v", err)
	} else {
		return nil
	}
}

func unwatchDirByEvent(t *fileTailer, event inotifyEvent) {
	watchedDirsAfter := make([]*Dir, 0, len(t.watchedDirs)-1)
	for _, existing := range t.watchedDirs {
		if existing.wd != int(event.Wd) {
			watchedDirsAfter = append(watchedDirsAfter, existing)
		}
	}
	t.watchedDirs = watchedDirsAfter
}

func (w *watcher) runFseventProducerLoop() fseventProducerLoop {
	return runInotifyLoop(w.fd)
}

func initWatcher() (fswatcher, Error) {
	fd, err := syscall.InotifyInit1(syscall.IN_CLOEXEC)
	if err != nil {
		return nil, NewError(NotSpecified, err, "inotify_init1() failed")
	}
	return &watcher{fd: fd}, nil
}

func (w *watcher) watchDir(path string) (*Dir, Error) {
	var (
		dir *Dir
		err error
		Err Error
	)
	dir, Err = newDir(path)
	if Err != nil {
		return nil, Err
	}
	dir.wd, err = syscall.InotifyAddWatch(w.fd, path, syscall.IN_MODIFY|syscall.IN_MOVED_FROM|syscall.IN_MOVED_TO|syscall.IN_DELETE|syscall.IN_CREATE)
	if err != nil {
		return nil, NewErrorf(NotSpecified, err, "%q: inotify_add_watch() failed", path)
	}
	return dir, nil
}

func newDir(path string) (*Dir, Error) {
	return &Dir{path: path}, nil
}

func (w *watcher) watchFile(_ fileMeta) Error {
	// nothing to do, because on Linux we watch the directory and don't need to watch individual files.
	return nil
}

func findDir(t *fileTailer, event inotifyEvent) (*Dir, Error) {
	for _, dir := range t.watchedDirs {
		if dir.wd == int(event.Wd) {
			return dir, nil
		}
	}
	// We should always find the directory in watchedDirs if we receive an inotifyEvent.
	// The following code should never be executed.
	// However, Github issue #124 is a bug report where we receive an inotifyEvent for an unwatched directory.
	// Create a detailed error message to help figuring out why this happens.
	var (
		errMsg strings.Builder
		first  = true
	)
	fmt.Fprintf(&errMsg, "watch list inconsistent: cannot find directory for file system event %q with watch descriptor %d in the list of watched directories [", event.String(), event.Wd)
	for _, dir := range t.watchedDirs {
		if !first {
			fmt.Fprintf(&errMsg, ", ")
		}
		fmt.Fprintf(&errMsg, "Wd=%d (%q)", dir.wd, dir.Path())
		first = false
	}
	fmt.Fprintf(&errMsg, "]")
	return nil, NewError(NotSpecified, nil, errMsg.String())
}

func (w *watcher) processEvent(t *fileTailer, fsevent fsevent, log logrus.FieldLogger) Error {
	event, ok := fsevent.(inotifyEvent)
	if !ok {
		return NewErrorf(NotSpecified, nil, "received a file system event of unknown type %T", event)
	}
	dir, Err := findDir(t, event)
	if Err != nil {
		return Err
	}
	dirLogger := log.WithField("directory", dir.path)
	dirLogger.Debugf("received event: %v", event)
	if event.Mask&syscall.IN_IGNORED == syscall.IN_IGNORED {
		unwatchDirByEvent(t, event) // need to remove it from watchedDirs, because otherwise we close the removed dir on shutdown which causes an error
		return NewErrorf(NotSpecified, nil, "%s: directory was removed while being watched", dir.path)
	}
	if event.Mask&syscall.IN_MODIFY == syscall.IN_MODIFY {
		file, ok := t.watchedFiles[filepath.Join(dir.path, event.Name)]
		if !ok {
			return nil // unrelated file was modified
		}
		truncated, err := isTruncated(file.file)
		if err != nil {
			return NewErrorf(NotSpecified, err, "%v: seek() or stat() failed", file.file.Name())
		}
		if truncated {
			_, err = file.file.Seek(0, io.SeekStart)
			if err != nil {
				return NewErrorf(NotSpecified, err, "%v: seek() failed", file.file.Name())
			}
			file.reader.Clear()
		}
		readErr := t.readNewLines(file, dirLogger)
		if readErr != nil {
			return readErr
		}
	}
	if event.Mask&syscall.IN_MOVED_FROM == syscall.IN_MOVED_FROM || event.Mask&syscall.IN_DELETE == syscall.IN_DELETE || event.Mask&syscall.IN_CREATE == syscall.IN_CREATE || event.Mask&syscall.IN_MOVED_TO == syscall.IN_MOVED_TO {
		// There are a lot of corner cases here:
		// * a file is renamed, but still matches the pattern so we continue watching it (MOVED_FROM followed by MOVED_TO)
		// * a file is created overwriting an existing file
		// * a file is moved to the watched directory overwriting an existing file
		// Trying to figure out what happened from the events would be error prone.
		// Therefore, we don't care which of the above events we received, we just update our watched files with the current
		// state of the watched directory.
		err := t.syncFilesInDir(dir, true, dirLogger)
		if err != nil {
			return err
		}
	}
	return nil
}

func isTruncated(file *os.File) (bool, error) {
	currentPos, err := file.Seek(0, io.SeekCurrent)
	if err != nil {
		return false, err
	}
	fileInfo, err := file.Stat()
	if err != nil {
		return false, err
	}
	return currentPos > fileInfo.Size(), nil
}

func findSameFile(t *fileTailer, file os.FileInfo, _ string) (*fileWithReader, Error) {
	var (
		fileInfo os.FileInfo
		err      error
	)
	for _, watchedFile := range t.watchedFiles {
		fileInfo, err = watchedFile.file.Stat()
		if err != nil {
			return nil, NewErrorf(NotSpecified, err, "%v: stat failed", watchedFile.file.Name())
		}
		if os.SameFile(fileInfo, file) {
			return watchedFile, nil
		}
	}
	return nil, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/route53"
)

func resourceAwsRoute53Record() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsRoute53RecordCreate,
		Read:   resourceAwsRoute53RecordRead,
		Delete: resourceAwsRoute53RecordDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"type": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"zone_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ttl": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"records": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
				ForceNew: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},
		},
	}
}

func resourceAwsRoute53RecordCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).route53

	zone := d.Get("zone_id").(string)

	zoneRecord, err := conn.GetHostedZone(zone)
	if err != nil {
		return err
	}

	// Check if the current record name contains the zone suffix.
	// If it does not, add the zone name to form a fully qualified name
	// and keep AWS happy.
	recordName := d.Get("name").(string)
	zoneName := strings.Trim(zoneRecord.HostedZone.Name, ".")
	if ok := strings.HasSuffix(recordName, zoneName); !ok {
		d.Set("name", strings.Join([]string{recordName, zoneName}, "."))
	}

	// Get the record
	rec, err := resourceAwsRoute53RecordBuildSet(d)
	if err != nil {
		return err
	}

	// Create the new records. We abuse StateChangeConf for this to
	// retry for us since Route53 sometimes returns errors about another
	// operation happening at the same time.
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Managed by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "UPSERT",
				Record: *rec,
			},
		},
	}

	log.Printf("[DEBUG] Creating resource records for zone: %s, name: %s",
		zone, d.Get("name").(string))

	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			resp, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return nil, "rejected", nil
				}

				return nil, "failure", err
			}

			return resp.ChangeInfo, "accepted", nil
		},
	}

	respRaw, err := wait.WaitForState()
	if err != nil {
		return err
	}
	changeInfo := respRaw.(route53.ChangeInfo)

	// Generate an ID
	d.SetId(fmt.Sprintf("%s_%s_%s", zone, d.Get("name").(string), d.Get("type").(string)))

	// Wait until we are done
	wait = resource.StateChangeConf{
		Delay:      30 * time.Second,
		Pending:    []string{"PENDING"},
		Target:     "INSYNC",
		Timeout:    10 * time.Minute,
		MinTimeout: 5 * time.Second,
		Refresh: func() (result interface{}, state string, err error) {
			return resourceAwsRoute53Wait(conn, changeInfo.ID)
		},
	}
	_, err = wait.WaitForState()
	if err != nil {
		return err
	}

	return nil
}

func resourceAwsRoute53RecordRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).route53

	zone := d.Get("zone_id").(string)
	lopts := &route53.ListOpts{
		Name: d.Get("name").(string),
		Type: d.Get("type").(string),
	}
	resp, err := conn.ListResourceRecordSets(zone, lopts)
	if err != nil {
		return err
	}

	// Scan for a matching record
	found := false
	for _, record := range resp.Records {
		if route53.FQDN(record.Name) != route53.FQDN(lopts.Name) {
			continue
		}
		if strings.ToUpper(record.Type) != strings.ToUpper(lopts.Type) {
			continue
		}

		found = true

		d.Set("records", record.Records)
		d.Set("ttl", record.TTL)

		break
	}

	if !found {
		d.SetId("")
	}

	return nil
}

func resourceAwsRoute53RecordDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).route53

	// Get the records
	rec, err := resourceAwsRoute53RecordBuildSet(d)
	if err != nil {
		return err
	}

	// Create the new records
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Deleted by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "DELETE",
				Record: *rec,
			},
		},
	}
	zone := d.Get("zone_id").(string)
	log.Printf("[DEBUG] Deleting resource records for zone: %s, name: %s",
		zone, d.Get("name").(string))

	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			_, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return 42, "rejected", nil
				}

				if strings.Contains(err.Error(), "InvalidChangeBatch") {
					// This means that the record is already gone.
					return 42, "accepted", nil
				}

				return 42, "failure", err
			}

			return 42, "accepted", nil
		},
	}

	if _, err := wait.WaitForState(); err != nil {
		return err
	}

	return nil
}

func resourceAwsRoute53RecordBuildSet(d *schema.ResourceData) (*route53.ResourceRecordSet, error) {
	recs := d.Get("records").(*schema.Set).List()
	records := make([]string, 0, len(recs))

	for _, r := range recs {
		records = append(records, r.(string))
	}

	rec := &route53.ResourceRecordSet{
		Name:    d.Get("name").(string),
		Type:    d.Get("type").(string),
		TTL:     d.Get("ttl").(int),
		Records: records,
	}
	return rec, nil
}
<commit_msg>simplify HasSuffix check<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/route53"
)

func resourceAwsRoute53Record() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsRoute53RecordCreate,
		Read:   resourceAwsRoute53RecordRead,
		Delete: resourceAwsRoute53RecordDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"type": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"zone_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ttl": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"records": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
				ForceNew: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},
		},
	}
}

func resourceAwsRoute53RecordCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).route53

	zone := d.Get("zone_id").(string)

	zoneRecord, err := conn.GetHostedZone(zone)
	if err != nil {
		return err
	}

	// Check if the current record name contains the zone suffix.
	// If it does not, add the zone name to form a fully qualified name
	// and keep AWS happy.
	recordName := d.Get("name").(string)
	zoneName := strings.Trim(zoneRecord.HostedZone.Name, ".")
	if !strings.HasSuffix(recordName, zoneName) {
		d.Set("name", strings.Join([]string{recordName, zoneName}, "."))
	}

	// Get the record
	rec, err := resourceAwsRoute53RecordBuildSet(d)
	if err != nil {
		return err
	}

	// Create the new records. We abuse StateChangeConf for this to
	// retry for us since Route53 sometimes returns errors about another
	// operation happening at the same time.
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Managed by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "UPSERT",
				Record: *rec,
			},
		},
	}

	log.Printf("[DEBUG] Creating resource records for zone: %s, name: %s",
		zone, d.Get("name").(string))

	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			resp, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return nil, "rejected", nil
				}

				return nil, "failure", err
			}

			return resp.ChangeInfo, "accepted", nil
		},
	}

	respRaw, err := wait.WaitForState()
	if err != nil {
		return err
	}
	changeInfo := respRaw.(route53.ChangeInfo)

	// Generate an ID
	d.SetId(fmt.Sprintf("%s_%s_%s", zone, d.Get("name").(string), d.Get("type").(string)))

	// Wait until we are done
	wait = resource.StateChangeConf{
		Delay:      30 * time.Second,
		Pending:    []string{"PENDING"},
		Target:     "INSYNC",
		Timeout:    10 * time.Minute,
		MinTimeout: 5 * time.Second,
		Refresh: func() (result interface{}, state string, err error) {
			return resourceAwsRoute53Wait(conn, changeInfo.ID)
		},
	}
	_, err = wait.WaitForState()
	if err != nil {
		return err
	}

	return nil
}

func resourceAwsRoute53RecordRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).route53

	zone := d.Get("zone_id").(string)
	lopts := &route53.ListOpts{
		Name: d.Get("name").(string),
		Type: d.Get("type").(string),
	}
	resp, err := conn.ListResourceRecordSets(zone, lopts)
	if err != nil {
		return err
	}

	// Scan for a matching record
	found := false
	for _, record := range resp.Records {
		if route53.FQDN(record.Name) != route53.FQDN(lopts.Name) {
			continue
		}
		if strings.ToUpper(record.Type) != strings.ToUpper(lopts.Type) {
			continue
		}

		found = true

		d.Set("records", record.Records)
		d.Set("ttl", record.TTL)

		break
	}

	if !found {
		d.SetId("")
	}

	return nil
}

func resourceAwsRoute53RecordDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).route53

	// Get the records
	rec, err := resourceAwsRoute53RecordBuildSet(d)
	if err != nil {
		return err
	}

	// Create the new records
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Deleted by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "DELETE",
				Record: *rec,
			},
		},
	}
	zone := d.Get("zone_id").(string)
	log.Printf("[DEBUG] Deleting resource records for zone: %s, name: %s",
		zone, d.Get("name").(string))

	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			_, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return 42, "rejected", nil
				}

				if strings.Contains(err.Error(), "InvalidChangeBatch") {
					// This means that the record is already gone.
					return 42, "accepted", nil
				}

				return 42, "failure", err
			}

			return 42, "accepted", nil
		},
	}

	if _, err := wait.WaitForState(); err != nil {
		return err
	}

	return nil
}

func resourceAwsRoute53RecordBuildSet(d *schema.ResourceData) (*route53.ResourceRecordSet, error) {
	recs := d.Get("records").(*schema.Set).List()
	records := make([]string, 0, len(recs))

	for _, r := range recs {
		records = append(records, r.(string))
	}

	rec := &route53.ResourceRecordSet{
		Name:    d.Get("name").(string),
		Type:    d.Get("type").(string),
		TTL:     d.Get("ttl").(int),
		Records: records,
	}
	return rec, nil
}
<|endoftext|>