<commit_before>package fasthttp

import (
	"bytes"
	"errors"
	"fmt"
	"html"
	"io"
	"mime"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

// FSHandlerCacheDuration is the duration for caching open file handles
// by FSHandler.
const FSHandlerCacheDuration = 5 * time.Second

// FSHandler returns request handler serving static files from
// the given root folder.
//
// stripSlashes indicates how many leading slashes must be stripped
// from requested path before searching requested file in the root folder.
// Examples:
//
//   * stripSlashes = 0, original path: "/foo/bar", result: "/foo/bar"
//   * stripSlashes = 1, original path: "/foo/bar", result: "/bar"
//   * stripSlashes = 2, original path: "/foo/bar", result: ""
//
// FSHandler caches requested file handles for FSHandlerCacheDuration.
// Make sure your program has enough 'max open files' limit aka
// 'ulimit -n' if root folder contains many files.
//
// Do not create multiple FSHandler instances for the same (root, stripSlashes)
// arguments - just reuse a single instance. Otherwise goroutine leak
// will occur.
func FSHandler(root string, stripSlashes int) RequestHandler {
	// strip trailing slashes from the root path
	for len(root) > 0 && root[len(root)-1] == '/' {
		root = root[:len(root)-1]
	}

	// serve files from the current working directory if root is empty
	if len(root) == 0 {
		root = "."
	}

	if stripSlashes < 0 {
		stripSlashes = 0
	}

	h := &fsHandler{
		root:         root,
		stripSlashes: stripSlashes,
		cache:        make(map[string]*fsFile),
	}
	go func() {
		for {
			time.Sleep(FSHandlerCacheDuration / 2)
			h.cleanCache()
		}
	}()
	return h.handleRequest
}

type fsHandler struct {
	root         string
	stripSlashes int
	cache        map[string]*fsFile
	pendingFiles []*fsFile
	cacheLock    sync.Mutex

	fileReaderPool sync.Pool
}

type fsFile struct {
	h             *fsHandler
	f             *os.File
	dirIndex      []byte
	contentType   string
	contentLength int

	t            time.Time
	readersCount int
}

func (ff *fsFile) Reader(incrementReaders bool) io.Reader {
	if incrementReaders {
		ff.h.cacheLock.Lock()
		ff.readersCount++
		ff.h.cacheLock.Unlock()
	}

	v := ff.h.fileReaderPool.Get()
	if v == nil {
		r := &fsFileReader{
			ff: ff,
		}
		r.v = r
		return r
	}
	r := v.(*fsFileReader)
	r.ff = ff
	if r.offset > 0 {
		panic("BUG: fsFileReader with non-nil offset found in the pool")
	}

	return r
}

func (ff *fsFile) Release() {
	if ff.f != nil {
		ff.f.Close()
	}
}

type fsFileReader struct {
	ff     *fsFile
	offset int64

	v interface{}
}

func (r *fsFileReader) Close() error {
	ff := r.ff

	ff.h.cacheLock.Lock()
	ff.readersCount--
	if ff.readersCount < 0 {
		panic("BUG: negative fsFile.readersCount!")
	}
	ff.h.cacheLock.Unlock()

	r.ff = nil
	r.offset = 0
	ff.h.fileReaderPool.Put(r.v)
	return nil
}

func (r *fsFileReader) Read(p []byte) (int, error) {
	ff := r.ff

	if ff.f != nil {
		n, err := ff.f.ReadAt(p, r.offset)
		r.offset += int64(n)
		return n, err
	}

	if r.offset == int64(len(ff.dirIndex)) {
		return 0, io.EOF
	}
	n := copy(p, ff.dirIndex[r.offset:])
	r.offset += int64(n)
	return n, nil
}

func (r *fsFileReader) WriteTo(w io.Writer) (int64, error) {
	if r.offset != 0 {
		panic("BUG: WriteTo must not be called after Read")
	}

	ff := r.ff

	var err error
	if ff.f != nil {
		pv := copyBufPool.Get()
		p := pv.([]byte)
		for err == nil {
			n, err := ff.f.ReadAt(p, r.offset)
			r.offset += int64(n)
		}
		copyBufPool.Put(pv)
		if err == io.EOF {
			err = nil
		}
	} else {
		var n int
		n, err = w.Write(ff.dirIndex)
		r.offset = int64(n)
	}
	return r.offset, err
}

func (h *fsHandler) cleanCache() {
	t := time.Now()
	h.cacheLock.Lock()

	// Close files which couldn't be closed before due to non-zero
	// readers count.
	var pendingFiles []*fsFile
	for _, ff := range h.pendingFiles {
		if ff.readersCount > 0 {
			pendingFiles = append(pendingFiles, ff)
		} else {
			ff.Release()
		}
	}
	h.pendingFiles = pendingFiles

	// Close stale file handles.
	for k, ff := range h.cache {
		if t.Sub(ff.t) > FSHandlerCacheDuration {
			if ff.readersCount > 0 {
				// There are pending readers on stale file handle,
				// so we cannot close it. Put it into pendingFiles
				// so it will be closed later.
				h.pendingFiles = append(h.pendingFiles, ff)
			} else {
				ff.Release()
			}
			delete(h.cache, k)
		}
	}

	h.cacheLock.Unlock()
}

func (h *fsHandler) handleRequest(ctx *RequestCtx) {
	path := ctx.Path()
	path = stripPathSlashes(path, h.stripSlashes)

	if n := bytes.IndexByte(path, 0); n >= 0 {
		ctx.Logger().Printf("cannot serve path with nil byte at position %d: %q", n, path)
		ctx.Error("Are you a hacker?", StatusBadRequest)
		return
	}

	incrementReaders := true

	h.cacheLock.Lock()
	ff, ok := h.cache[string(path)]
	if ok {
		ff.readersCount++
		incrementReaders = false
	}
	h.cacheLock.Unlock()

	if !ok {
		pathStr := string(path)
		filePath := h.root + pathStr
		var err error
		ff, err = h.openFSFile(filePath)
		if err == errDirIndexRequired {
			ff, err = h.createDirIndex(ctx.URI(), filePath)
			if err != nil {
				ctx.Logger().Printf("Cannot create index for directory %q: %s", filePath, err)
				ctx.Error("Cannot create directory index", StatusNotFound)
				return
			}
		} else if err != nil {
			ctx.Logger().Printf("cannot open file %q: %s", filePath, err)
			ctx.Error("Cannot open requested path", StatusNotFound)
			return
		}

		h.cacheLock.Lock()
		ff1, ok := h.cache[pathStr]
		if !ok {
			h.cache[pathStr] = ff
		}
		h.cacheLock.Unlock()

		if ok {
			// The file has been already opened by another
			// goroutine, so close the current file and use
			// the file opened by another goroutine instead.
			ff.Release()
			ff = ff1
		}
	}

	ctx.SetBodyStream(ff.Reader(incrementReaders), ff.contentLength)
	ctx.SetContentType(ff.contentType)
}

var errDirIndexRequired = errors.New("directory index required")

func (h *fsHandler) createDirIndex(base *URI, filePath string) (*fsFile, error) {
	var buf bytes.Buffer
	w := &buf

	basePathEscaped := html.EscapeString(string(base.Path()))
	fmt.Fprintf(w, "<title>%s</title>", basePathEscaped)
	fmt.Fprintf(w, "
%s
", basePathEscaped)
	fmt.Fprintf(w, "
")

	if len(basePathEscaped) > 1 {
		fmt.Fprintf(w, `<li><a href="..">..</a></li>`)
	}

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

	filenames, err := f.Readdirnames(0)
	f.Close()
	if err != nil {
		return nil, err
	}

	var u URI
	base.CopyTo(&u)

	sort.Sort(sort.StringSlice(filenames))
	for _, name := range filenames {
		u.Update(name)
		pathEscaped := html.EscapeString(string(u.Path()))
		fmt.Fprintf(w, `<li><a href="%s">%s</a></li>`, pathEscaped, html.EscapeString(name))
	}

	fmt.Fprintf(w, "
")
	dirIndex := w.Bytes()

	ff := &fsFile{
		h:             h,
		dirIndex:      dirIndex,
		contentType:   "text/html",
		contentLength: len(dirIndex),
	}
	return ff, nil
}

func (h *fsHandler) openFSFile(filePath string) (*fsFile, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}

	stat, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, err
	}

	if stat.IsDir() {
		f.Close()

		indexPath := filePath + "/index.html"
		ff, err := h.openFSFile(indexPath)
		if err == nil {
			return ff, nil
		}
		if !os.IsNotExist(err) {
			return nil, err
		}
		return nil, errDirIndexRequired
	}

	n := stat.Size()
	contentLength := int(n)
	if n != int64(contentLength) {
		f.Close()
		return nil, fmt.Errorf("too big file: %d bytes", n)
	}

	ext := fileExtension(filePath)
	contentType := mime.TypeByExtension(ext)

	ff := &fsFile{
		h:             h,
		f:             f,
		contentType:   contentType,
		contentLength: contentLength,
	}
	return ff, nil
}

func stripPathSlashes(path []byte, stripSlashes int) []byte {
	// strip leading slashes
	for stripSlashes > 0 && len(path) > 0 {
		if path[0] != '/' {
			panic("BUG: path must start with slash")
		}
		n := bytes.IndexByte(path[1:], '/')
		if n < 0 {
			path = path[:0]
			break
		}
		path = path[n+1:]
		stripSlashes--
	}

	// strip trailing slashes
	for len(path) > 0 && path[len(path)-1] == '/' {
		path = path[:len(path)-1]
	}

	return path
}

func fileExtension(path string) string {
	n := strings.LastIndexByte(path, '.')
	if n < 0 {
		return ""
	}
	return path[n:]
}
<commit_msg>typo fix<commit_after>package fasthttp

import (
	"bytes"
	"errors"
	"fmt"
	"html"
	"io"
	"mime"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

// FSHandlerCacheDuration is the duration for caching open file handles
// by FSHandler.
const FSHandlerCacheDuration = 5 * time.Second

// FSHandler returns request handler serving static files from
// the given root folder.
//
// stripSlashes indicates how many leading slashes must be stripped
// from requested path before searching requested file in the root folder.
// Examples:
//
//   * stripSlashes = 0, original path: "/foo/bar", result: "/foo/bar"
//   * stripSlashes = 1, original path: "/foo/bar", result: "/bar"
//   * stripSlashes = 2, original path: "/foo/bar", result: ""
//
// FSHandler caches requested file handles for FSHandlerCacheDuration.
// Make sure your program has enough 'max open files' limit aka
// 'ulimit -n' if root folder contains many files.
//
// Do not create multiple FSHandler instances for the same (root, stripSlashes)
// arguments - just reuse a single instance. Otherwise goroutine leak
// will occur.
func FSHandler(root string, stripSlashes int) RequestHandler {
	// strip trailing slashes from the root path
	for len(root) > 0 && root[len(root)-1] == '/' {
		root = root[:len(root)-1]
	}

	// serve files from the current working directory if root is empty
	if len(root) == 0 {
		root = "."
	}

	if stripSlashes < 0 {
		stripSlashes = 0
	}

	h := &fsHandler{
		root:         root,
		stripSlashes: stripSlashes,
		cache:        make(map[string]*fsFile),
	}
	go func() {
		for {
			time.Sleep(FSHandlerCacheDuration / 2)
			h.cleanCache()
		}
	}()
	return h.handleRequest
}

type fsHandler struct {
	root         string
	stripSlashes int
	cache        map[string]*fsFile
	pendingFiles []*fsFile
	cacheLock    sync.Mutex

	fileReaderPool sync.Pool
}

type fsFile struct {
	h             *fsHandler
	f             *os.File
	dirIndex      []byte
	contentType   string
	contentLength int

	t            time.Time
	readersCount int
}

func (ff *fsFile) Reader(incrementReaders bool) io.Reader {
	if incrementReaders {
		ff.h.cacheLock.Lock()
		ff.readersCount++
		ff.h.cacheLock.Unlock()
	}

	v := ff.h.fileReaderPool.Get()
	if v == nil {
		r := &fsFileReader{
			ff: ff,
		}
		r.v = r
		return r
	}
	r := v.(*fsFileReader)
	r.ff = ff
	if r.offset > 0 {
		panic("BUG: fsFileReader with non-nil offset found in the pool")
	}

	return r
}

func (ff *fsFile) Release() {
	if ff.f != nil {
		ff.f.Close()
	}
}

type fsFileReader struct {
	ff     *fsFile
	offset int64

	v interface{}
}

func (r *fsFileReader) Close() error {
	ff := r.ff

	ff.h.cacheLock.Lock()
	ff.readersCount--
	if ff.readersCount < 0 {
		panic("BUG: negative fsFile.readersCount!")
	}
	ff.h.cacheLock.Unlock()

	r.ff = nil
	r.offset = 0
	ff.h.fileReaderPool.Put(r.v)
	return nil
}

func (r *fsFileReader) Read(p []byte) (int, error) {
	ff := r.ff

	if ff.f != nil {
		n, err := ff.f.ReadAt(p, r.offset)
		r.offset += int64(n)
		return n, err
	}

	if r.offset == int64(len(ff.dirIndex)) {
		return 0, io.EOF
	}
	n := copy(p, ff.dirIndex[r.offset:])
	r.offset += int64(n)
	return n, nil
}

func (r *fsFileReader) WriteTo(w io.Writer) (int64, error) {
	if r.offset != 0 {
		panic("BUG: WriteTo must not be called after Read")
	}

	ff := r.ff

	var err error
	var n int
	if ff.f != nil {
		pv := copyBufPool.Get()
		p := pv.([]byte)
		for err == nil {
			n, err = ff.f.ReadAt(p, r.offset)
			r.offset += int64(n)
		}
		copyBufPool.Put(pv)
		if err == io.EOF {
			err = nil
		}
	} else {
		n, err = w.Write(ff.dirIndex)
		r.offset = int64(n)
	}
	return r.offset, err
}

func (h *fsHandler) cleanCache() {
	t := time.Now()
	h.cacheLock.Lock()

	// Close files which couldn't be closed before due to non-zero
	// readers count.
	var pendingFiles []*fsFile
	for _, ff := range h.pendingFiles {
		if ff.readersCount > 0 {
			pendingFiles = append(pendingFiles, ff)
		} else {
			ff.Release()
		}
	}
	h.pendingFiles = pendingFiles

	// Close stale file handles.
	for k, ff := range h.cache {
		if t.Sub(ff.t) > FSHandlerCacheDuration {
			if ff.readersCount > 0 {
				// There are pending readers on stale file handle,
				// so we cannot close it. Put it into pendingFiles
				// so it will be closed later.
				h.pendingFiles = append(h.pendingFiles, ff)
			} else {
				ff.Release()
			}
			delete(h.cache, k)
		}
	}

	h.cacheLock.Unlock()
}

func (h *fsHandler) handleRequest(ctx *RequestCtx) {
	path := ctx.Path()
	path = stripPathSlashes(path, h.stripSlashes)

	if n := bytes.IndexByte(path, 0); n >= 0 {
		ctx.Logger().Printf("cannot serve path with nil byte at position %d: %q", n, path)
		ctx.Error("Are you a hacker?", StatusBadRequest)
		return
	}

	incrementReaders := true

	h.cacheLock.Lock()
	ff, ok := h.cache[string(path)]
	if ok {
		ff.readersCount++
		incrementReaders = false
	}
	h.cacheLock.Unlock()

	if !ok {
		pathStr := string(path)
		filePath := h.root + pathStr
		var err error
		ff, err = h.openFSFile(filePath)
		if err == errDirIndexRequired {
			ff, err = h.createDirIndex(ctx.URI(), filePath)
			if err != nil {
				ctx.Logger().Printf("Cannot create index for directory %q: %s", filePath, err)
				ctx.Error("Cannot create directory index", StatusNotFound)
				return
			}
		} else if err != nil {
			ctx.Logger().Printf("cannot open file %q: %s", filePath, err)
			ctx.Error("Cannot open requested path", StatusNotFound)
			return
		}

		h.cacheLock.Lock()
		ff1, ok := h.cache[pathStr]
		if !ok {
			h.cache[pathStr] = ff
		}
		h.cacheLock.Unlock()

		if ok {
			// The file has been already opened by another
			// goroutine, so close the current file and use
			// the file opened by another goroutine instead.
			ff.Release()
			ff = ff1
		}
	}

	ctx.SetBodyStream(ff.Reader(incrementReaders), ff.contentLength)
	ctx.SetContentType(ff.contentType)
}

var errDirIndexRequired = errors.New("directory index required")

func (h *fsHandler) createDirIndex(base *URI, filePath string) (*fsFile, error) {
	var buf bytes.Buffer
	w := &buf

	basePathEscaped := html.EscapeString(string(base.Path()))
	fmt.Fprintf(w, "<title>%s</title>", basePathEscaped)
	fmt.Fprintf(w, "
%s
", basePathEscaped)
	fmt.Fprintf(w, "
")

	if len(basePathEscaped) > 1 {
		fmt.Fprintf(w, `<li><a href="..">..</a></li>`)
	}

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

	filenames, err := f.Readdirnames(0)
	f.Close()
	if err != nil {
		return nil, err
	}

	var u URI
	base.CopyTo(&u)

	sort.Sort(sort.StringSlice(filenames))
	for _, name := range filenames {
		u.Update(name)
		pathEscaped := html.EscapeString(string(u.Path()))
		fmt.Fprintf(w, `<li><a href="%s">%s</a></li>`, pathEscaped, html.EscapeString(name))
	}

	fmt.Fprintf(w, "
")
	dirIndex := w.Bytes()

	ff := &fsFile{
		h:             h,
		dirIndex:      dirIndex,
		contentType:   "text/html",
		contentLength: len(dirIndex),
	}
	return ff, nil
}

func (h *fsHandler) openFSFile(filePath string) (*fsFile, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}

	stat, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, err
	}

	if stat.IsDir() {
		f.Close()

		indexPath := filePath + "/index.html"
		ff, err := h.openFSFile(indexPath)
		if err == nil {
			return ff, nil
		}
		if !os.IsNotExist(err) {
			return nil, err
		}
		return nil, errDirIndexRequired
	}

	n := stat.Size()
	contentLength := int(n)
	if n != int64(contentLength) {
		f.Close()
		return nil, fmt.Errorf("too big file: %d bytes", n)
	}

	ext := fileExtension(filePath)
	contentType := mime.TypeByExtension(ext)

	ff := &fsFile{
		h:             h,
		f:             f,
		contentType:   contentType,
		contentLength: contentLength,
	}
	return ff, nil
}

func stripPathSlashes(path []byte, stripSlashes int) []byte {
	// strip leading slashes
	for stripSlashes > 0 && len(path) > 0 {
		if path[0] != '/' {
			panic("BUG: path must start with slash")
		}
		n := bytes.IndexByte(path[1:], '/')
		if n < 0 {
			path = path[:0]
			break
		}
		path = path[n+1:]
		stripSlashes--
	}

	// strip trailing slashes
	for len(path) > 0 && path[len(path)-1] == '/' {
		path = path[:len(path)-1]
	}

	return path
}

func fileExtension(path string) string {
	n := strings.LastIndexByte(path, '.')
	if n < 0 {
		return ""
	}
	return path[n:]
}
<|endoftext|>


<commit_before>package dedup_test

import (
	"bytes"
	"io"
	"testing"

	"io/ioutil"

	"github.com/klauspost/dedup"
)

func TestReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
	blocks := r.BlockSizes()
	for _, s := range blocks[:len(blocks)-1] {
		if s != size {
			t.Fatal("wrong size, expected", size, "got", s)
		}
	}
}

func TestReaderStream(t *testing.T) {
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewStreamWriter(&data, dedup.ModeFixed, size, 10*size)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewStreamReader(&data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestSeekReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 50<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%100)*size : (i%100)*size+size]
		// Write into the following ones
		dst := b[(100+i)*size : (i+100)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewSeekReader(&idx, bytes.NewReader(data.Bytes()))
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestDynamicRoundtrip(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeDynamic, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Dynamic Index size:", idx.Len())
	t.Log("Dynamic Data size:", data.Len())

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")
	blocks := r.BlockSizes()
	avg := 0
	for _, v := range blocks {
		if v > size {
			t.Fatal("too big block returned, should not be >", size, "was", v)
		}
		avg += v
	}
	t.Log("Average block size:", avg/len(blocks), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

//TODO: Test WriteTo

func BenchmarkReader64K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	_, err = io.Copy(w, input)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx = bytes.NewBuffer(index)
		data = bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil && err != io.EOF {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short, expected", len(b), "was", n)
		}
		err = r.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}

func BenchmarkReader4K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 4 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	_, err = io.Copy(w, input)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx := bytes.NewBuffer(index)
		data := bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil && err != io.EOF {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short, expected", len(b), "was", n)
		}
		err = r.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}

func BenchmarkReader1K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 1 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	_, err = io.Copy(w, input)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx := bytes.NewBuffer(index)
		data := bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil && err != io.EOF {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short, expected", len(b), "was", n)
		}
		err = r.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}
<commit_msg>Add "WriteTo" test.<commit_after>package dedup_test

import (
	"bytes"
	"io"
	"testing"

	"io/ioutil"

	"github.com/klauspost/dedup"
)

func TestReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
	blocks := r.BlockSizes()
	for _, s := range blocks[:len(blocks)-1] {
		if s != size {
			t.Fatal("wrong size, expected", size, "got", s)
		}
	}
}

func TestReaderStream(t *testing.T) {
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewStreamWriter(&data, dedup.ModeFixed, size, 10*size)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewStreamReader(&data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestSeekReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 50<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%100)*size : (i%100)*size+size]
		// Write into the following ones
		dst := b[(100+i)*size : (i+100)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewSeekReader(&idx, bytes.NewReader(data.Bytes()))
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestDynamicRoundtrip(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeDynamic, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Dynamic Index size:", idx.Len())
	t.Log("Dynamic Data size:", data.Len())

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")
	blocks := r.BlockSizes()
	avg := 0
	for _, v := range blocks {
		if v > size {
			t.Fatal("too big block returned, should not be >", size, "was", v)
		}
		avg += v
	}
	t.Log("Average block size:", avg/len(blocks), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestReaderWriteTo(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	dst := &bytes.Buffer{}
	n, err := r.WriteTo(dst)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != int(n) {
		t.Errorf("Write count, expected n %d, got %d", len(b), n)
	}

	out := dst.Bytes()
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
	blocks := r.BlockSizes()
	for _, s := range blocks[:len(blocks)-1] {
		if s != size {
			t.Fatal("wrong size, expected", size, "got", s)
		}
	}
}

func BenchmarkReader64K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	_, err = io.Copy(w, input)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx = bytes.NewBuffer(index)
		data = bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil && err != io.EOF {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short, expected", len(b), "was", n)
		}
		err = r.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}

func BenchmarkReader4K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 4 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	_, err = io.Copy(w, input)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx := bytes.NewBuffer(index)
		data := bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil && err != io.EOF {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short, expected", len(b), "was", n)
		}
		err = r.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}

func BenchmarkReader1K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 1 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	_, err = io.Copy(w, input)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx := bytes.NewBuffer(index)
		data := bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil && err != io.EOF {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short, expected", len(b), "was", n)
		}
		err = r.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>package dedup_test

import (
	"bytes"
	"io"
	"testing"

	"io/ioutil"

	"github.com/klauspost/dedup"
)

func TestReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
	blocks := r.BlockSizes()
	for _, s := range blocks[:len(blocks)-1] {
		if s != size {
			t.Fatal("wrong size, expected", size, "got", s)
		}
	}
}

func TestReaderStream(t *testing.T) {
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewStreamWriter(&data, dedup.ModeFixed, size, 10)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewStreamReader(&data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestSeekReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 50<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%100)*size : (i%100)*size+size]
		// Write into the following ones
		dst := b[(100+i)*size : (i+100)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewSeekReader(&idx, bytes.NewReader(data.Bytes()))
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestDynamicRoundtrip(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 30<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeDynamic, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Dynamic Index size:", idx.Len())
	t.Log("Dynamic Data size:", data.Len())

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")
	blocks := r.BlockSizes()
	avg := 0
	for _, v := range blocks {
		if v > size {
			t.Fatal("too big block returned, should not be >", size, "was", v)
		}
		avg += v
	}
	t.Log("Average block size:", avg/len(blocks), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func BenchmarkReader64K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx = bytes.NewBuffer(index)
		data = bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short")
		}
		err = w.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}

func BenchmarkReader4K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 4 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx = bytes.NewBuffer(index)
		data = bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short")
		}
		err = w.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}
<commit_msg>Decrease test size.<commit_after>package dedup_test

import (
	"bytes"
	"io"
	"testing"

	"io/ioutil"

	"github.com/klauspost/dedup"
)

func TestReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
	blocks := r.BlockSizes()
	for _, s := range blocks[:len(blocks)-1] {
		if s != size {
			t.Fatal("wrong size, expected", size, "got", s)
		}
	}
}

func TestReaderStream(t *testing.T) {
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewStreamWriter(&data, dedup.ModeFixed, size, 10)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewStreamReader(&data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestSeekReader(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 50<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%100)*size : (i%100)*size+size]
		// Write into the following ones
		dst := b[(100+i)*size : (i+100)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Fixed Index size:", idx.Len())
	t.Log("Fixed Data size:", data.Len(), "-", data.Len()*100/totalinput, "%")

	r, err := dedup.NewSeekReader(&idx, bytes.NewReader(data.Bytes()))
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestDynamicRoundtrip(t *testing.T) {
	idx := bytes.Buffer{}
	data := bytes.Buffer{}

	const totalinput = 10<<20 + 65
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(&idx, &data, dedup.ModeDynamic, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Dynamic Index size:", idx.Len())
	t.Log("Dynamic Data size:", data.Len())

	r, err := dedup.NewReader(&idx, &data)
	if err != nil {
		t.Fatal(err)
	}

	t.Log("Maximum estimated memory:", r.MaxMem(), "bytes")
	blocks := r.BlockSizes()
	avg := 0
	for _, v := range blocks {
		if v > size {
			t.Fatal("too big block returned, should not be >", size, "was", v)
		}
		avg += v
	}
	t.Log("Average block size:", avg/len(blocks), "bytes")

	out, err := ioutil.ReadAll(r)
	if err != io.EOF && err != nil {
		t.Fatal(err)
	}
	if len(b) != len(out) {
		t.Fatalf("Expected len %d, got %d", len(b), len(out))
	}
	if bytes.Compare(b, out) != 0 {
		t.Fatal("Output mismatch")
	}
	err = r.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func BenchmarkReader64K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 64 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 50; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx = bytes.NewBuffer(index)
		data = bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short")
		}
		err = w.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}

func BenchmarkReader4K(t *testing.B) {
	idx := &bytes.Buffer{}
	data := &bytes.Buffer{}

	const totalinput = 10 << 20
	input := getBufferSize(totalinput)

	const size = 4 << 10
	b := input.Bytes()
	// Create some duplicates
	for i := 0; i < 500; i++ {
		// Read from 10 first blocks
		src := b[(i%10)*size : (i%10)*size+size]
		// Write into the following ones
		dst := b[(10+i)*size : (i+10)*size+size]
		copy(dst, src)
	}
	input = bytes.NewBuffer(b)
	w, err := dedup.NewWriter(idx, data, dedup.ModeFixed, size, 0)
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(w, input)
	err = w.Close()
	if err != nil {
		t.Fatal(err)
	}

	index := idx.Bytes()
	alldata := data.Bytes()

	t.ResetTimer()
	t.SetBytes(totalinput)
	for i := 0; i < t.N; i++ {
		idx = bytes.NewBuffer(index)
		data = bytes.NewBuffer(alldata)
		r, err := dedup.NewReader(idx, data)
		if err != nil {
			t.Fatal(err)
		}
		n, err := io.Copy(ioutil.Discard, r)
		if err != nil {
			t.Fatal(err)
		}
		if n != int64(len(b)) {
			t.Fatal("read was short")
		}
		err = w.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>package generate

import (
	"go/ast"
	"go/parser"
	"go/token"
	"testing"
)

func TestGeneratedSourceParsesWithoutErrors(t *testing.T) {
	pkg := "pkg"
	sumType := "Sum"
	variants := []string{"X", "Y", "Z"}

	src := Generate(pkg, sumType, variants...)

	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "src.go", src, 0)
	if err != nil {
		t.Error(err)
	}
}

func TestGeneratedSourceDeclaresTypeWithNameGiven(t *testing.T) {
	pkg := "pkg"
	sumType := "Sum"
	variants := []string{"X", "Y", "Z"}

	src := Generate(pkg, sumType, variants...)

	fset := token.NewFileSet()
	f, _ := parser.ParseFile(fset, "src.go", src, 0)

	found := false
	for _, typ := range typesInFile(f) {
		if typ.Name.Name == sumType {
			found = true
		}
	}

	if !found {
		t.Errorf("type %s not declared in generated sources", sumType)
	}
}

func typesInFile(f *ast.File) []*ast.TypeSpec {
	var types []*ast.TypeSpec

	for _, decl := range f.Decls {
		decl := decl.(*ast.GenDecl)

		if decl.Tok == token.TYPE {

			for _, spec := range decl.Specs {
				if typ, ok := spec.(*ast.TypeSpec); ok {

					types = append(types, typ)
				}
			}
		}
	}

	return types
}

func TestGeneratedSourceHasRightPackageClause(t *testing.T) {
	for _, pkg := range []string{"pkg", "pack"} {
		sumType := "Sum"
		variants := []string{"X", "Y", "Z"}

		src := Generate(pkg, sumType, variants...)

		fset := token.NewFileSet()
		f, _ := parser.ParseFile(fset, "src.go", src, parser.PackageClauseOnly)

		if f.Name.Name != pkg {
			t.Errorf("generated file should belong to package %s not %s",
				pkg, f.Name.Name)
		}
	}
}
<commit_msg>Add TestGeneratedSumTypeIsInterface<commit_after>package generate

import (
	"go/ast"
	"go/parser"
	"go/token"
	"testing"
)

func TestGeneratedSourceParsesWithoutErrors(t *testing.T) {
	pkg := "pkg"
	sumType := "Sum"
	variants := []string{"X", "Y", "Z"}

	src := Generate(pkg, sumType, variants...)

	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "src.go", src, 0)
	if err != nil {
		t.Error(err)
	}
}

func TestGeneratedSourceDeclaresTypeWithNameGiven(t *testing.T) {
	pkg := "pkg"
	sumType := "Sum"
	variants := []string{"X", "Y", "Z"}

	src := Generate(pkg, sumType, variants...)

	fset := token.NewFileSet()
	f, _ := parser.ParseFile(fset, "src.go", src, 0)

	found := false
	for _, typ := range typesInFile(f) {
		if typ.Name.Name == sumType {
			found = true
		}
	}

	if !found {
		t.Errorf("type %s not declared in generated sources", sumType)
	}
}

func typesInFile(f *ast.File) []*ast.TypeSpec {
	var types []*ast.TypeSpec

	for _, decl := range f.Decls {
		decl := decl.(*ast.GenDecl)

		if decl.Tok == token.TYPE {

			for _, spec := range decl.Specs {
				if typ, ok := spec.(*ast.TypeSpec); ok {

					types = append(types, typ)
				}
			}
		}
	}

	return types
}

func TestGeneratedSourceHasRightPackageClause(t *testing.T) {
	for _, pkg := range []string{"pkg", "pack"} {
		sumType := "Sum"
		variants := []string{"X", "Y", "Z"}

		src := Generate(pkg, sumType, variants...)

		fset := token.NewFileSet()
		f, _ := parser.ParseFile(fset, "src.go", src, parser.PackageClauseOnly)

		if f.Name.Name != pkg {
			t.Errorf("generated file should belong to package %s not %s",
				pkg, f.Name.Name)
		}
	}
}

func TestGeneratedSumTypeIsInterface(t *testing.T) {
	pkg := "pkg"
	sumType := "Sum"
	variants := []string{"X", "Y", "Z"}

	src := Generate(pkg, sumType, variants...)

	fset := token.NewFileSet()
	f, _ := parser.ParseFile(fset, "src.go", src, 0)

	typ, ok := typeNamed(f, sumType)
	if !ok {
		t.Fatalf("generated source must contain type declaration for %s type", sumType)
	}

	_, isInterface := typ.Type.(*ast.InterfaceType)
	if !isInterface {
		t.Errorf("generated %s type should be an interface not a %T", sumType, typ.Type)
	}
}

// Looks for the declaration of a type named name in the specified file. The
// second return value is false if a type with the given name is not declared
// in the file.
func typeNamed(f *ast.File, name string) (*ast.TypeSpec, bool) {
	for _, typ := range typesInFile(f) {
		if typ.Name.Name == name {
			return typ, true
		}
	}
	return nil, false
}
<|endoftext|>


<commit_before>package msi

// #cgo CFLAGS: -I/usr/include/irods
// #cgo CXXFLAGS: -I/usr/include/irods -I/opt/irods-externals/boost1.60.0-0/include -I/opt/irods-externals/clang3.8-0/include/c++/v1 -nostdinc++ -std=c++14
// #cgo LDFLAGS: -Wl,-rpath,"/opt/irods-externals/boost1.60.0-0/lib" -Wl,-rpath,"/opt/irods-externals/clang3.8-0/lib" -lirods_server -lirods_common -lpthread -lc++ -lc++abi
/*
#include <stdlib.h>
#include "call_microservice.h"
*/
import "C"

import (
	"fmt"
	"unsafe"
)

var rei unsafe.Pointer

// Configure accepts a ruleExecInfo_t*, cast as an unsafe.Pointer.
// This is a requirement for passing C types between packages.
// You must run this function prior to executing microservices with msi.Call
func Configure(ruleExecInfo unsafe.Pointer) {
	rei = ruleExecInfo
}

// Call invokes a microservice by name. The variadic arguments can be *msi.Param,
// string, int, or int64.
func Call(msiName string, params ...interface{}) error {
	if rei == nil {
		return fmt.Errorf("Unable to call %v, ruleExecInfo is nil, please set using msi.Configure", msiName)
	}

	numParams := C.int(len(params))

	cParams := C.NewParamList(numParams)
	defer C.free(unsafe.Pointer(cParams))

	for inx, param := range params {
		var msParam *Param

		switch p := param.(type) {
		case string:
			msParam = NewParam(STR_MS_T)

			msParam.SetString(p)
		case int:
			msParam = NewParam(INT_MS_T)

			msParam.SetInt(p)
		case int64:
			msParam = NewParam(INT_MS_T)

			msParam.SetInt(int(p))
		case *Param:
			msParam = p
		case nil:
			msParam = new(Param)
		default:
			return fmt.Errorf("Parameter with unknown type passed in params")
		}

		C.SetMsParamListItem(cParams, C.int(inx), msParam.ptr)
	}

	var callInfo C.msiCallInfo_t
	size := unsafe.Sizeof(callInfo)
	C.bzero(unsafe.Pointer(&callInfo), C.size_t(size))

	cMsiName := C.CString(msiName)
	defer C.free(unsafe.Pointer(cMsiName))

	callInfo.microserviceName = cMsiName
	callInfo.params = cParams
	callInfo.paramsLen = numParams
	callInfo.rei = rei

	var errStr *C.char
	if status := C.call_microservice(&callInfo, &errStr); status < 0 {
		return fmt.Errorf("Error in call_microservice: %v", C.GoString(errStr))
	}

	return nil

}
<commit_msg>Updated documentation<commit_after>// Package msi is package that contains utilities for use in microservices written in Golang.
// GoRODS/msi provides a binding to the msParam_t type, and its various subtypes. In addition,
// it provides an interface for calling other microservices.
package msi

// #cgo CFLAGS: -I/usr/include/irods
// #cgo CXXFLAGS: -I/usr/include/irods -I/opt/irods-externals/boost1.60.0-0/include -I/opt/irods-externals/clang3.8-0/include/c++/v1 -nostdinc++ -std=c++14
// #cgo LDFLAGS: -Wl,-rpath,"/opt/irods-externals/boost1.60.0-0/lib" -Wl,-rpath,"/opt/irods-externals/clang3.8-0/lib" -lirods_server -lirods_common -lpthread -lc++ -lc++abi
/*
#include <stdlib.h>
#include "call_microservice.h"
*/
import "C"

import (
	"fmt"
	"unsafe"
)

var rei unsafe.Pointer

// Configure accepts a ruleExecInfo_t*, cast as an unsafe.Pointer.
// This is a requirement for passing C types between packages.
// You must run this function prior to executing microservices with msi.Call
func Configure(ruleExecInfo unsafe.Pointer) {
	rei = ruleExecInfo
}

// Call invokes a microservice by name. The variadic arguments can be *msi.Param,
// string, int, or int64.
func Call(msiName string, params ...interface{}) error {
	if rei == nil {
		return fmt.Errorf("Unable to call %v, ruleExecInfo is nil, please set using msi.Configure", msiName)
	}

	numParams := C.int(len(params))

	cParams := C.NewParamList(numParams)
	defer C.free(unsafe.Pointer(cParams))

	for inx, param := range params {
		var msParam *Param

		switch p := param.(type) {
		case string:
			msParam = NewParam(STR_MS_T)

			msParam.SetString(p)
		case int:
			msParam = NewParam(INT_MS_T)

			msParam.SetInt(p)
		case int64:
			msParam = NewParam(INT_MS_T)

			msParam.SetInt(int(p))
		case *Param:
			msParam = p
		case nil:
			msParam = new(Param)
		default:
			return fmt.Errorf("Parameter with unknown type passed in params")
		}

		C.SetMsParamListItem(cParams, C.int(inx), msParam.ptr)
	}

	var callInfo C.msiCallInfo_t
	size := unsafe.Sizeof(callInfo)
	C.bzero(unsafe.Pointer(&callInfo), C.size_t(size))

	cMsiName := C.CString(msiName)
	defer C.free(unsafe.Pointer(cMsiName))

	callInfo.microserviceName = cMsiName
	callInfo.params = cParams
	callInfo.paramsLen = numParams
	callInfo.rei = rei

	var errStr *C.char
	if status := C.call_microservice(&callInfo, &errStr); status < 0 {
		return fmt.Errorf("Error in call_microservice: %v", C.GoString(errStr))
	}

	return nil

}
<|endoftext|>


<commit_before>package matching

import (
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	. "github.com/SpectoLabs/hoverfly/core/util"
	"github.com/SpectoLabs/hoverfly/core/models"
	"github.com/SpectoLabs/hoverfly/core/views"
	"github.com/ryanuber/go-glob"
	"strings"
)

type RequestTemplateStore []RequestTemplateResponsePair

type RequestTemplateResponsePair struct {
	RequestTemplate RequestTemplate        `json:"requestTemplate"`
	Response        models.ResponseDetails `json:"response"`
}

type RequestTemplateResponsePairView struct {
	RequestTemplate RequestTemplate           `json:"requestTemplate"`
	Response        views.ResponseDetailsView `json:"response"`
}

type RequestTemplateResponsePairPayload struct {
	Data *[]RequestTemplateResponsePairView `json:"data"`
}

type RequestTemplate struct {
	Path        *string             `json:"path"`
	Method      *string             `json:"method"`
	Destination *string             `json:"destination"`
	Scheme      *string             `json:"scheme"`
	Query       *string             `json:"query"`
	Body        *string             `json:"body"`
	Headers     map[string][]string `json:"headers"`
}

func (this *RequestTemplateStore) GetResponse(req models.RequestDetails, webserver bool) (*models.ResponseDetails, error) {
	// iterate through the request templates, looking for template to match request
	for _, entry := range *this {
		// TODO: not matching by default on URL and body - need to enable this
		// TODO: need to enable regex matches
		// TODO: enable matching on scheme

		if entry.RequestTemplate.Body != nil && !glob.Glob(*entry.RequestTemplate.Body, req.Body) {
			continue
		}

		if !webserver {
			if entry.RequestTemplate.Destination != nil && !glob.Glob(*entry.RequestTemplate.Destination, req.Destination) {
				continue
			}
		}
		if entry.RequestTemplate.Path != nil && !glob.Glob(*entry.RequestTemplate.Path, req.Path) {
			continue
		}
		if entry.RequestTemplate.Query != nil && !glob.Glob(*entry.RequestTemplate.Query, req.Query) {
			continue
		}
		if !headerMatch(entry.RequestTemplate.Headers, req.Headers) {
			continue
		}
		if entry.RequestTemplate.Method != nil && !glob.Glob(*entry.RequestTemplate.Method, req.Method) {
			continue
		}

		// return the first template to match
		return &entry.Response, nil
	}
	return nil, errors.New("No match found")
}

// ImportPayloads - a function to save given payloads into the database.
func (this *RequestTemplateStore) ImportPayloads(pairPayload RequestTemplateResponsePairPayload) error {
	if len(*pairPayload.Data) > 0 {
		// Convert PayloadView back to Payload for internal storage
		templateStore := pairPayload.ConvertToRequestTemplateStore()
		for _, pl := range templateStore {

			//TODO: add hooks for concsistency with request import
			// note that importing hoverfly is a disallowed circular import

			*this = append(*this, pl)
		}
		log.WithFields(log.Fields{
			"total": len(*this),
		}).Info("payloads imported")
		return nil
	}
	return fmt.Errorf("Bad request. Nothing to import!")
}

func (this *RequestTemplateStore) Wipe() {
	// don't change the pointer here!
	*this = RequestTemplateStore{}
}

/**
Check keys and corresponding values in template headers are also present in request headers
*/
func headerMatch(templateHeaders, requestHeaders map[string][]string) bool {

	for templateHeaderKey, templateHeaderValues := range templateHeaders {
		for requestHeaderKey, requestHeaderValues := range requestHeaders {
			delete(requestHeaders, requestHeaderKey)
			requestHeaders[strings.ToLower(requestHeaderKey)] = requestHeaderValues

		}
		requestTemplateValues, ok := requestHeaders[strings.ToLower(templateHeaderKey)]
		if !ok {
			return false
		}

		for _, templateHeaderValue := range templateHeaderValues {
			found := 0
			for _, requestHeaderValue := range requestTemplateValues {
				if glob.Glob(strings.ToLower(templateHeaderValue), strings.ToLower(requestHeaderValue)) {
					found++
				}
			}

			if found == 0 {
				return false
			}
		}
	}
	return true
}

func (this *RequestTemplateStore) GetPayload() RequestTemplateResponsePairPayload {
	var pairsPayload []RequestTemplateResponsePairView
	for _, pair := range *this {
		pairsPayload = append(pairsPayload, pair.ConvertToRequestTemplateResponsePairView())
	}
	return RequestTemplateResponsePairPayload{
		Data: &pairsPayload,
	}
}

func (this *RequestTemplateResponsePair) ConvertToRequestTemplateResponsePairView() RequestTemplateResponsePairView {
	return RequestTemplateResponsePairView{
		RequestTemplate: this.RequestTemplate,
		Response:        this.Response.ConvertToResponseDetailsView(),
	}
}

func (this *RequestTemplateResponsePair) ConvertToRequestResponsePairView() views.RequestResponsePairView {

	return views.RequestResponsePairView{
		Request: views.RequestDetailsView{
			RequestType: StringToPointer("template"),
			Path:        this.RequestTemplate.Path,
			Method:      this.RequestTemplate.Method,
			Destination: this.RequestTemplate.Destination,
			Scheme:      this.RequestTemplate.Scheme,
			Query:       this.RequestTemplate.Query,
			Body:        this.RequestTemplate.Body,
			Headers:     this.RequestTemplate.Headers,
		},
		Response: this.Response.ConvertToResponseDetailsView(),
	}
}

func (this *RequestTemplateResponsePairPayload) ConvertToRequestTemplateStore() RequestTemplateStore {
	var requestTemplateStore RequestTemplateStore
	for _, pair := range *this.Data {
		requestTemplateStore = append(requestTemplateStore, pair.ConvertToRequestTemplateResponsePair())
	}
	return requestTemplateStore
}

func (this *RequestTemplateResponsePairView) ConvertToRequestTemplateResponsePair() RequestTemplateResponsePair {
	return RequestTemplateResponsePair{
		RequestTemplate: this.RequestTemplate,
		Response:        models.NewResponseDetailsFromResponseDetailsView(this.Response),
	}
}
<commit_msg>Renamed some variables in the template header matching to make it more clear<commit_after>package matching

import (
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	. "github.com/SpectoLabs/hoverfly/core/util"
	"github.com/SpectoLabs/hoverfly/core/models"
	"github.com/SpectoLabs/hoverfly/core/views"
	"github.com/ryanuber/go-glob"
	"strings"
)

type RequestTemplateStore []RequestTemplateResponsePair

type RequestTemplateResponsePair struct {
	RequestTemplate RequestTemplate        `json:"requestTemplate"`
	Response        models.ResponseDetails `json:"response"`
}

type RequestTemplateResponsePairView struct {
	RequestTemplate RequestTemplate           `json:"requestTemplate"`
	Response        views.ResponseDetailsView `json:"response"`
}

type RequestTemplateResponsePairPayload struct {
	Data *[]RequestTemplateResponsePairView `json:"data"`
}

type RequestTemplate struct {
	Path        *string             `json:"path"`
	Method      *string             `json:"method"`
	Destination *string             `json:"destination"`
	Scheme      *string             `json:"scheme"`
	Query       *string             `json:"query"`
	Body        *string             `json:"body"`
	Headers     map[string][]string `json:"headers"`
}

func (this *RequestTemplateStore) GetResponse(req models.RequestDetails, webserver bool) (*models.ResponseDetails, error) {
	// iterate through the request templates, looking for template to match request
	for _, entry := range *this {
		// TODO: not matching by default on URL and body - need to enable this
		// TODO: need to enable regex matches
		// TODO: enable matching on scheme

		if entry.RequestTemplate.Body != nil && !glob.Glob(*entry.RequestTemplate.Body, req.Body) {
			continue
		}

		if !webserver {
			if entry.RequestTemplate.Destination != nil && !glob.Glob(*entry.RequestTemplate.Destination, req.Destination) {
				continue
			}
		}
		if entry.RequestTemplate.Path != nil && !glob.Glob(*entry.RequestTemplate.Path, req.Path) {
			continue
		}
		if entry.RequestTemplate.Query != nil && !glob.Glob(*entry.RequestTemplate.Query, req.Query) {
			continue
		}
		if !headerMatch(entry.RequestTemplate.Headers, req.Headers) {
			continue
		}
		if entry.RequestTemplate.Method != nil && !glob.Glob(*entry.RequestTemplate.Method, req.Method) {
			continue
		}

		// return the first template to match
		return &entry.Response, nil
	}
	return nil, errors.New("No match found")
}

// ImportPayloads - a function to save given payloads into the database.
func (this *RequestTemplateStore) ImportPayloads(pairPayload RequestTemplateResponsePairPayload) error {
	if len(*pairPayload.Data) > 0 {
		// Convert PayloadView back to Payload for internal storage
		templateStore := pairPayload.ConvertToRequestTemplateStore()
		for _, pl := range templateStore {

			//TODO: add hooks for concsistency with request import
			// note that importing hoverfly is a disallowed circular import

			*this = append(*this, pl)
		}
		log.WithFields(log.Fields{
			"total": len(*this),
		}).Info("payloads imported")
		return nil
	}
	return fmt.Errorf("Bad request. Nothing to import!")
}

func (this *RequestTemplateStore) Wipe() {
	// don't change the pointer here!
	*this = RequestTemplateStore{}
}

/**
Check keys and corresponding values in template headers are also present in request headers
*/
func headerMatch(templateHeaders, requestHeaders map[string][]string) bool {

	for templateHeaderKey, templateHeaderValues := range templateHeaders {
		for requestHeaderKey, requestHeaderValues := range requestHeaders {
			delete(requestHeaders, requestHeaderKey)
			requestHeaders[strings.ToLower(requestHeaderKey)] = requestHeaderValues

		}
		requestTemplateValues, templateHeaderMatched := requestHeaders[strings.ToLower(templateHeaderKey)]
		if !templateHeaderMatched {
			return false
		}

		for _, templateHeaderValue := range templateHeaderValues {
			templateValueMatched := false
			for _, requestHeaderValue := range requestTemplateValues {
				if glob.Glob(strings.ToLower(templateHeaderValue), strings.ToLower(requestHeaderValue)) {
					templateValueMatched = true
				}
			}

			if !templateValueMatched {
				return false
			}
		}
	}
	return true
}

func (this *RequestTemplateStore) GetPayload() RequestTemplateResponsePairPayload {
	var pairsPayload []RequestTemplateResponsePairView
	for _, pair := range *this {
		pairsPayload = append(pairsPayload, pair.ConvertToRequestTemplateResponsePairView())
	}
	return RequestTemplateResponsePairPayload{
		Data: &pairsPayload,
	}
}

func (this *RequestTemplateResponsePair) ConvertToRequestTemplateResponsePairView() RequestTemplateResponsePairView {
	return RequestTemplateResponsePairView{
		RequestTemplate: this.RequestTemplate,
		Response:        this.Response.ConvertToResponseDetailsView(),
	}
}

func (this *RequestTemplateResponsePair) ConvertToRequestResponsePairView() views.RequestResponsePairView {

	return views.RequestResponsePairView{
		Request: views.RequestDetailsView{
			RequestType: StringToPointer("template"),
			Path:        this.RequestTemplate.Path,
			Method:      this.RequestTemplate.Method,
			Destination: this.RequestTemplate.Destination,
			Scheme:      this.RequestTemplate.Scheme,
			Query:       this.RequestTemplate.Query,
			Body:        this.RequestTemplate.Body,
			Headers:     this.RequestTemplate.Headers,
		},
		Response: this.Response.ConvertToResponseDetailsView(),
	}
}

func (this *RequestTemplateResponsePairPayload) ConvertToRequestTemplateStore() RequestTemplateStore {
	var requestTemplateStore RequestTemplateStore
	for _, pair := range *this.Data {
		requestTemplateStore = append(requestTemplateStore, pair.ConvertToRequestTemplateResponsePair())
	}
	return requestTemplateStore
}

func (this *RequestTemplateResponsePairView) ConvertToRequestTemplateResponsePair() RequestTemplateResponsePair {
	return RequestTemplateResponsePair{
		RequestTemplate: this.RequestTemplate,
		Response:        models.NewResponseDetailsFromResponseDetailsView(this.Response),
	}
}
<|endoftext|>


<commit_before>package veneur

import (
	"container/ring"
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	lightstep "github.com/lightstep/lightstep-tracer-go"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/sirupsen/logrus"
	"github.com/stripe/veneur/ssf"
	"github.com/stripe/veneur/trace"
)

const datadogResourceKey = "resource"
const datadogNameKey = "name"

const lightStepOperationKey = "name"

const totalSpansFlushedMetricKey = "worker.spans_flushed_total"

// spanSink is a receiver of spans that handles sending those spans to some
// downstream sink. Calls to `Ingest(span)` are meant to give the sink control
// of the span, with periodic calls to flush as a signal for sinks that don't
// handle their own flushing in a separate goroutine, etc.
type spanSink interface {
	Name() string
	Start(*trace.Client) error
	Ingest(ssf.SSFSpan) error
	Flush()
}

// DatadogSpanSink is a sink for sending spans to a Datadog trace agent.
type datadogSpanSink struct {
	HTTPClient   *http.Client
	buffer       *ring.Ring
	bufferSize   int
	mutex        *sync.Mutex
	stats        *statsd.Client
	commonTags   map[string]string
	traceAddress string
	traceClient  *trace.Client
}

// NewDatadogSpanSink creates a new Datadog sink for trace spans.
func NewDatadogSpanSink(config *Config, stats *statsd.Client, httpClient *http.Client, commonTags map[string]string) (*datadogSpanSink, error) {
	return &datadogSpanSink{
		HTTPClient:   httpClient,
		bufferSize:   config.SsfBufferSize,
		buffer:       ring.New(config.SsfBufferSize),
		mutex:        &sync.Mutex{},
		stats:        stats,
		commonTags:   commonTags,
		traceAddress: config.DatadogTraceAPIAddress,
	}, nil
}

// Name returns the name of this sink.
func (dd *datadogSpanSink) Name() string {
	return "datadog"
}

// Start performs final adjustments on the sink.
func (dd *datadogSpanSink) Start(cl *trace.Client) error {
	dd.traceClient = cl
	return nil
}

// Ingest takes the span and adds it to the ringbuffer.
func (dd *datadogSpanSink) Ingest(span ssf.SSFSpan) error {
	dd.mutex.Lock()
	defer dd.mutex.Unlock()

	dd.buffer.Value = span
	dd.buffer = dd.buffer.Next()
	return nil
}

// Flush signals the sink to send it's spans to their destination. For this
// sync it means we'll be making an HTTP request to send them along. We assume
// it's beneficial to performance to defer these until the normal 10s flush.
func (dd *datadogSpanSink) Flush() {
	dd.mutex.Lock()

	ssfSpans := make([]ssf.SSFSpan, 0, dd.buffer.Len())

	dd.buffer.Do(func(t interface{}) {
		const tooEarly = 1497
		const tooLate = 1497629343000000

		if t != nil {
			ssfSpan, ok := t.(ssf.SSFSpan)
			if !ok {
				log.Error("Got an unknown object in tracing ring!")
				// We'll just skip this one so we don't poison pill or anything.
				return
			}

			var timeErr string
			if ssfSpan.StartTimestamp < tooEarly {
				timeErr = "type:tooEarly"
			}
			if ssfSpan.StartTimestamp > tooLate {
				timeErr = "type:tooLate"
			}
			if timeErr != "" {
				dd.stats.Incr("worker.trace.sink.timestamp_error", []string{timeErr}, 1) // TODO tag as dd?
			}

			if ssfSpan.Tags == nil {
				ssfSpan.Tags = make(map[string]string)
			}

			// Add common tags from veneur's config
			// this will overwrite tags already present on the span
			for k, v := range dd.commonTags {
				ssfSpan.Tags[k] = v
			}
			ssfSpans = append(ssfSpans, ssfSpan)
		}
	})

	// Reset the ring.
	dd.buffer = ring.New(dd.bufferSize)

	// We're done manipulating stuff, let Ingest loose again.
	dd.mutex.Unlock()

	serviceCount := make(map[string]int64)
	var finalTraces []*DatadogTraceSpan
	// Conver the SSFSpans into Datadog Spans
	for _, span := range ssfSpans {
		// -1 is a canonical way of passing in invalid info in Go
		// so we should support that too
		parentID := span.ParentId

		// check if this is the root span
		if parentID <= 0 {
			// we need parentId to be zero for json:omitempty to work
			parentID = 0
		}

		resource := span.Tags[datadogResourceKey]
		name := span.Name

		tags := map[string]string{}
		// Get the span's existing tags
		for k, v := range span.Tags {
			tags[k] = v
		}

		delete(tags, datadogResourceKey)

		// TODO implement additional metrics
		var metrics map[string]float64

		var errorCode int64
		if span.Error {
			errorCode = 2
		}

		ddspan := &DatadogTraceSpan{
			TraceID:  span.TraceId,
			SpanID:   span.Id,
			ParentID: parentID,
			Service:  span.Service,
			Name:     name,
			Resource: resource,
			Start:    span.StartTimestamp,
			Duration: span.EndTimestamp - span.StartTimestamp,
			// TODO don't hardcode
			Type:    "http",
			Error:   errorCode,
			Metrics: metrics,
			Meta:    tags,
		}
		serviceCount[span.Service]++
		finalTraces = append(finalTraces, ddspan)
	}

	if len(finalTraces) != 0 {
		// this endpoint is not documented to take an array... but it does
		// another curious constraint of this endpoint is that it does not
		// support "Content-Encoding: deflate"

		err := postHelper(context.TODO(), dd.HTTPClient, dd.stats, dd.traceClient, fmt.Sprintf("%s/spans", dd.traceAddress), finalTraces, "flush_traces", false)

		if err == nil {
			log.WithField("traces", len(finalTraces)).Info("Completed flushing traces to Datadog")
			dd.stats.Count(totalSpansFlushedMetricKey, int64(len(ssfSpans)), []string{"sink:datadog"}, 1)
			// TODO: Per service counters?
		} else {
			log.WithFields(logrus.Fields{
				"traces":        len(finalTraces),
				logrus.ErrorKey: err}).Warn("Error flushing traces to Datadog")
		}
		for service, count := range serviceCount {
			dd.stats.Count(totalSpansFlushedMetricKey, count, []string{"sink:datadog", fmt.Sprintf("service:%s", service)}, 1)
		}
	} else {
		log.Info("No traces to flush to Datadog, skipping.")
	}
}

// lightStepSpanSink is a sink for spans to be sent to the LightStep client.
type lightStepSpanSink struct {
	tracers      []opentracing.Tracer
	stats        *statsd.Client
	commonTags   map[string]string
	mutex        *sync.Mutex
	serviceCount map[string]int64
	traceClient  *trace.Client
}

// NewLightStepSpanSink creates a new instance of a LightStepSpanSink.
func NewLightStepSpanSink(config *Config, stats *statsd.Client, commonTags map[string]string) (*lightStepSpanSink, error) {

	var host *url.URL
	host, err := url.Parse(config.TraceLightstepCollectorHost)
	if err != nil {
		log.WithError(err).WithField(
			"host", config.TraceLightstepCollectorHost,
		).Error("Error parsing LightStep collector URL")
		return &lightStepSpanSink{}, err
	}

	port, err := strconv.Atoi(host.Port())
	if err != nil {
		log.WithError(err).WithFields(logrus.Fields{
			"port":         port,
			"default_port": lightstepDefaultPort,
		}).Warn("Error parsing LightStep port, using default")
		port = lightstepDefaultPort
	}

	reconPeriod := lightstepDefaultInterval
	if config.TraceLightstepReconnectPeriod != "" {
		reconPeriod, err = time.ParseDuration(config.TraceLightstepReconnectPeriod)
		if err != nil {
			log.WithError(err).WithFields(logrus.Fields{
				"interval":         config.TraceLightstepReconnectPeriod,
				"default_interval": lightstepDefaultInterval,
			}).Warn("Failed to parse reconnect duration, using default.")
			reconPeriod = lightstepDefaultInterval
		}
	}

	log.WithFields(logrus.Fields{
		"Host": host.Hostname(),
		"Port": port,
	}).Info("Dialing lightstep host")

	maxSpans := config.TraceLightstepMaximumSpans
	if maxSpans == 0 {
		maxSpans = config.SsfBufferSize
		log.WithField("max spans", maxSpans).Info("Using default maximum spans — ssf_buffer_size — for LightStep")
	}

	lightstepMultiplexTracerNum := config.TraceLightstepNumClients
	// If config value is missing, this value should default to one client
	if lightstepMultiplexTracerNum <= 0 {
		lightstepMultiplexTracerNum = 1
	}

	tracers := make([]opentracing.Tracer, 0, lightstepMultiplexTracerNum)

	for i := 0; i < lightstepMultiplexTracerNum; i++ {
		tracers = append(tracers, lightstep.NewTracer(lightstep.Options{
			AccessToken:     config.TraceLightstepAccessToken,
			ReconnectPeriod: reconPeriod,
			Collector: lightstep.Endpoint{
				Host:      host.Hostname(),
				Port:      port,
				Plaintext: true,
			},
			UseGRPC:          true,
			MaxBufferedSpans: maxSpans,
		}))
	}

	return &lightStepSpanSink{
		tracers:      tracers,
		stats:        stats,
		serviceCount: make(map[string]int64),
		mutex:        &sync.Mutex{},
	}, nil
}

func (ls *lightStepSpanSink) Start(cl *trace.Client) error {
	ls.traceClient = cl
	return nil
}

// Name returns this sink's name.
func (ls *lightStepSpanSink) Name() string {
	return "lightstep"
}

// Ingest takes in a span and passed it along to the LS client after
// some sanity checks and improvements are made.
func (ls *lightStepSpanSink) Ingest(ssfSpan ssf.SSFSpan) error {
	parentID := ssfSpan.ParentId
	if parentID <= 0 {
		parentID = 0
	}

	var errorCode int64
	if ssfSpan.Error {
		errorCode = 1
	}

	timestamp := time.Unix(ssfSpan.StartTimestamp/1e9, ssfSpan.StartTimestamp%1e9)

	if len(ls.tracers) == 0 {
		err := fmt.Errorf("No lightstep tracer clients initialized")
		log.Error(err)
		return err
	}
	// pick the tracer to use
	tracerIndex := ssfSpan.TraceId % int64(len(ls.tracers))
	tracer := ls.tracers[tracerIndex]

	sp := tracer.StartSpan(
		ssfSpan.Name,
		opentracing.StartTime(timestamp),
		lightstep.SetTraceID(uint64(ssfSpan.TraceId)),
		lightstep.SetSpanID(uint64(ssfSpan.Id)),
		lightstep.SetParentSpanID(uint64(parentID)))

	sp.SetTag(trace.ResourceKey, ssfSpan.Tags[trace.ResourceKey]) // TODO Why is this here?
	sp.SetTag(lightstep.ComponentNameKey, ssfSpan.Service)
	// TODO don't hardcode
	sp.SetTag("type", "http")
	sp.SetTag("error-code", errorCode)
	for k, v := range ssfSpan.Tags {
		sp.SetTag(k, v)
	}
	// And now set any veneur common tags
	for k, v := range ls.commonTags {
		sp.SetTag(k, v)
	}
	// TODO add metrics as tags to the span as well?

	if errorCode > 0 {
		// Note: this sets the OT-standard "error" tag, which
		// LightStep uses to flag error spans.
		ext.Error.Set(sp, true)
	}

	endTime := time.Unix(ssfSpan.EndTimestamp/1e9, ssfSpan.EndTimestamp%1e9)
	finishOpts := opentracing.FinishOptions{FinishTime: endTime}
	if sp, ok := sp.(*trace.Span); ok {
		sp.ClientFinishWithOptions(ls.traceClient, finishOpts)
	} else {
		sp.FinishWithOptions(finishOpts)
	}

	service := ssfSpan.Service
	if service == "" {
		service = "unknown"
	}

	// Protect mutating the service count with a mutex
	ls.mutex.Lock()
	defer ls.mutex.Unlock()
	ls.serviceCount[service]++

	return nil
}

// Flush doesn't need to do anything to the LS tracer, so we emit metrics
// instead.
func (ls *lightStepSpanSink) Flush() {
	ls.mutex.Lock()
	defer ls.mutex.Unlock()

	totalCount := int64(0)
	for service, count := range ls.serviceCount {
		totalCount += count
		ls.stats.Count(totalSpansFlushedMetricKey, count, []string{"sink:lightstep", fmt.Sprintf("service:%s", service)}, 1)
	}
	ls.serviceCount = make(map[string]int64)
	log.WithField("total_spans", totalCount).Debug("Checkpointing flushed spans for Lightstep")
}
<commit_msg>Make LS Span ingestion feed the constructed span into the LS tracer<commit_after>package veneur

import (
	"container/ring"
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	lightstep "github.com/lightstep/lightstep-tracer-go"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/sirupsen/logrus"
	"github.com/stripe/veneur/ssf"
	"github.com/stripe/veneur/trace"
)

const datadogResourceKey = "resource"
const datadogNameKey = "name"

const lightStepOperationKey = "name"

const totalSpansFlushedMetricKey = "worker.spans_flushed_total"

// spanSink is a receiver of spans that handles sending those spans to some
// downstream sink. Calls to `Ingest(span)` are meant to give the sink control
// of the span, with periodic calls to flush as a signal for sinks that don't
// handle their own flushing in a separate goroutine, etc.
type spanSink interface {
	Name() string
	Start(*trace.Client) error
	Ingest(ssf.SSFSpan) error
	Flush()
}

// DatadogSpanSink is a sink for sending spans to a Datadog trace agent.
type datadogSpanSink struct {
	HTTPClient   *http.Client
	buffer       *ring.Ring
	bufferSize   int
	mutex        *sync.Mutex
	stats        *statsd.Client
	commonTags   map[string]string
	traceAddress string
	traceClient  *trace.Client
}

// NewDatadogSpanSink creates a new Datadog sink for trace spans.
func NewDatadogSpanSink(config *Config, stats *statsd.Client, httpClient *http.Client, commonTags map[string]string) (*datadogSpanSink, error) {
	return &datadogSpanSink{
		HTTPClient:   httpClient,
		bufferSize:   config.SsfBufferSize,
		buffer:       ring.New(config.SsfBufferSize),
		mutex:        &sync.Mutex{},
		stats:        stats,
		commonTags:   commonTags,
		traceAddress: config.DatadogTraceAPIAddress,
	}, nil
}

// Name returns the name of this sink.
func (dd *datadogSpanSink) Name() string {
	return "datadog"
}

// Start performs final adjustments on the sink.
func (dd *datadogSpanSink) Start(cl *trace.Client) error {
	dd.traceClient = cl
	return nil
}

// Ingest takes the span and adds it to the ringbuffer.
func (dd *datadogSpanSink) Ingest(span ssf.SSFSpan) error {
	dd.mutex.Lock()
	defer dd.mutex.Unlock()

	dd.buffer.Value = span
	dd.buffer = dd.buffer.Next()
	return nil
}

// Flush signals the sink to send it's spans to their destination. For this
// sync it means we'll be making an HTTP request to send them along. We assume
// it's beneficial to performance to defer these until the normal 10s flush.
func (dd *datadogSpanSink) Flush() {
	dd.mutex.Lock()

	ssfSpans := make([]ssf.SSFSpan, 0, dd.buffer.Len())

	dd.buffer.Do(func(t interface{}) {
		const tooEarly = 1497
		const tooLate = 1497629343000000

		if t != nil {
			ssfSpan, ok := t.(ssf.SSFSpan)
			if !ok {
				log.Error("Got an unknown object in tracing ring!")
				// We'll just skip this one so we don't poison pill or anything.
				return
			}

			var timeErr string
			if ssfSpan.StartTimestamp < tooEarly {
				timeErr = "type:tooEarly"
			}
			if ssfSpan.StartTimestamp > tooLate {
				timeErr = "type:tooLate"
			}
			if timeErr != "" {
				dd.stats.Incr("worker.trace.sink.timestamp_error", []string{timeErr}, 1) // TODO tag as dd?
			}

			if ssfSpan.Tags == nil {
				ssfSpan.Tags = make(map[string]string)
			}

			// Add common tags from veneur's config
			// this will overwrite tags already present on the span
			for k, v := range dd.commonTags {
				ssfSpan.Tags[k] = v
			}
			ssfSpans = append(ssfSpans, ssfSpan)
		}
	})

	// Reset the ring.
	dd.buffer = ring.New(dd.bufferSize)

	// We're done manipulating stuff, let Ingest loose again.
	dd.mutex.Unlock()

	serviceCount := make(map[string]int64)
	var finalTraces []*DatadogTraceSpan
	// Conver the SSFSpans into Datadog Spans
	for _, span := range ssfSpans {
		// -1 is a canonical way of passing in invalid info in Go
		// so we should support that too
		parentID := span.ParentId

		// check if this is the root span
		if parentID <= 0 {
			// we need parentId to be zero for json:omitempty to work
			parentID = 0
		}

		resource := span.Tags[datadogResourceKey]
		name := span.Name

		tags := map[string]string{}
		// Get the span's existing tags
		for k, v := range span.Tags {
			tags[k] = v
		}

		delete(tags, datadogResourceKey)

		// TODO implement additional metrics
		var metrics map[string]float64

		var errorCode int64
		if span.Error {
			errorCode = 2
		}

		ddspan := &DatadogTraceSpan{
			TraceID:  span.TraceId,
			SpanID:   span.Id,
			ParentID: parentID,
			Service:  span.Service,
			Name:     name,
			Resource: resource,
			Start:    span.StartTimestamp,
			Duration: span.EndTimestamp - span.StartTimestamp,
			// TODO don't hardcode
			Type:    "http",
			Error:   errorCode,
			Metrics: metrics,
			Meta:    tags,
		}
		serviceCount[span.Service]++
		finalTraces = append(finalTraces, ddspan)
	}

	if len(finalTraces) != 0 {
		// this endpoint is not documented to take an array... but it does
		// another curious constraint of this endpoint is that it does not
		// support "Content-Encoding: deflate"

		err := postHelper(context.TODO(), dd.HTTPClient, dd.stats, dd.traceClient, fmt.Sprintf("%s/spans", dd.traceAddress), finalTraces, "flush_traces", false)

		if err == nil {
			log.WithField("traces", len(finalTraces)).Info("Completed flushing traces to Datadog")
			dd.stats.Count(totalSpansFlushedMetricKey, int64(len(ssfSpans)), []string{"sink:datadog"}, 1)
			// TODO: Per service counters?
		} else {
			log.WithFields(logrus.Fields{
				"traces":        len(finalTraces),
				logrus.ErrorKey: err}).Warn("Error flushing traces to Datadog")
		}
		for service, count := range serviceCount {
			dd.stats.Count(totalSpansFlushedMetricKey, count, []string{"sink:datadog", fmt.Sprintf("service:%s", service)}, 1)
		}
	} else {
		log.Info("No traces to flush to Datadog, skipping.")
	}
}

// lightStepSpanSink is a sink for spans to be sent to the LightStep client.
type lightStepSpanSink struct {
	tracers      []opentracing.Tracer
	stats        *statsd.Client
	commonTags   map[string]string
	mutex        *sync.Mutex
	serviceCount map[string]int64
	traceClient  *trace.Client
}

// NewLightStepSpanSink creates a new instance of a LightStepSpanSink.
func NewLightStepSpanSink(config *Config, stats *statsd.Client, commonTags map[string]string) (*lightStepSpanSink, error) {

	var host *url.URL
	host, err := url.Parse(config.TraceLightstepCollectorHost)
	if err != nil {
		log.WithError(err).WithField(
			"host", config.TraceLightstepCollectorHost,
		).Error("Error parsing LightStep collector URL")
		return &lightStepSpanSink{}, err
	}

	port, err := strconv.Atoi(host.Port())
	if err != nil {
		log.WithError(err).WithFields(logrus.Fields{
			"port":         port,
			"default_port": lightstepDefaultPort,
		}).Warn("Error parsing LightStep port, using default")
		port = lightstepDefaultPort
	}

	reconPeriod := lightstepDefaultInterval
	if config.TraceLightstepReconnectPeriod != "" {
		reconPeriod, err = time.ParseDuration(config.TraceLightstepReconnectPeriod)
		if err != nil {
			log.WithError(err).WithFields(logrus.Fields{
				"interval":         config.TraceLightstepReconnectPeriod,
				"default_interval": lightstepDefaultInterval,
			}).Warn("Failed to parse reconnect duration, using default.")
			reconPeriod = lightstepDefaultInterval
		}
	}

	log.WithFields(logrus.Fields{
		"Host": host.Hostname(),
		"Port": port,
	}).Info("Dialing lightstep host")

	maxSpans := config.TraceLightstepMaximumSpans
	if maxSpans == 0 {
		maxSpans = config.SsfBufferSize
		log.WithField("max spans", maxSpans).Info("Using default maximum spans — ssf_buffer_size — for LightStep")
	}

	lightstepMultiplexTracerNum := config.TraceLightstepNumClients
	// If config value is missing, this value should default to one client
	if lightstepMultiplexTracerNum <= 0 {
		lightstepMultiplexTracerNum = 1
	}

	tracers := make([]opentracing.Tracer, 0, lightstepMultiplexTracerNum)

	for i := 0; i < lightstepMultiplexTracerNum; i++ {
		tracers = append(tracers, lightstep.NewTracer(lightstep.Options{
			AccessToken:     config.TraceLightstepAccessToken,
			ReconnectPeriod: reconPeriod,
			Collector: lightstep.Endpoint{
				Host:      host.Hostname(),
				Port:      port,
				Plaintext: true,
			},
			UseGRPC:          true,
			MaxBufferedSpans: maxSpans,
		}))
	}

	return &lightStepSpanSink{
		tracers:      tracers,
		stats:        stats,
		serviceCount: make(map[string]int64),
		mutex:        &sync.Mutex{},
	}, nil
}

func (ls *lightStepSpanSink) Start(cl *trace.Client) error {
	ls.traceClient = cl
	return nil
}

// Name returns this sink's name.
func (ls *lightStepSpanSink) Name() string {
	return "lightstep"
}

// Ingest takes in a span and passed it along to the LS client after
// some sanity checks and improvements are made.
func (ls *lightStepSpanSink) Ingest(ssfSpan ssf.SSFSpan) error {
	parentID := ssfSpan.ParentId
	if parentID <= 0 {
		parentID = 0
	}

	var errorCode int64
	if ssfSpan.Error {
		errorCode = 1
	}

	timestamp := time.Unix(ssfSpan.StartTimestamp/1e9, ssfSpan.StartTimestamp%1e9)

	if len(ls.tracers) == 0 {
		err := fmt.Errorf("No lightstep tracer clients initialized")
		log.Error(err)
		return err
	}
	// pick the tracer to use
	tracerIndex := ssfSpan.TraceId % int64(len(ls.tracers))
	tracer := ls.tracers[tracerIndex]

	sp := tracer.StartSpan(
		ssfSpan.Name,
		opentracing.StartTime(timestamp),
		lightstep.SetTraceID(uint64(ssfSpan.TraceId)),
		lightstep.SetSpanID(uint64(ssfSpan.Id)),
		lightstep.SetParentSpanID(uint64(parentID)))

	sp.SetTag(trace.ResourceKey, ssfSpan.Tags[trace.ResourceKey]) // TODO Why is this here?
	sp.SetTag(lightstep.ComponentNameKey, ssfSpan.Service)
	// TODO don't hardcode
	sp.SetTag("type", "http")
	sp.SetTag("error-code", errorCode)
	for k, v := range ssfSpan.Tags {
		sp.SetTag(k, v)
	}
	// And now set any veneur common tags
	for k, v := range ls.commonTags {
		sp.SetTag(k, v)
	}
	// TODO add metrics as tags to the span as well?

	if errorCode > 0 {
		// Note: this sets the OT-standard "error" tag, which
		// LightStep uses to flag error spans.
		ext.Error.Set(sp, true)
	}

	endTime := time.Unix(ssfSpan.EndTimestamp/1e9, ssfSpan.EndTimestamp%1e9)
	finishOpts := opentracing.FinishOptions{FinishTime: endTime}
	sp.FinishWithOptions(finishOpts)

	service := ssfSpan.Service
	if service == "" {
		service = "unknown"
	}

	// Protect mutating the service count with a mutex
	ls.mutex.Lock()
	defer ls.mutex.Unlock()
	ls.serviceCount[service]++

	return nil
}

// Flush doesn't need to do anything to the LS tracer, so we emit metrics
// instead.
func (ls *lightStepSpanSink) Flush() {
	ls.mutex.Lock()
	defer ls.mutex.Unlock()

	totalCount := int64(0)
	for service, count := range ls.serviceCount {
		totalCount += count
		ls.stats.Count(totalSpansFlushedMetricKey, count, []string{"sink:lightstep", fmt.Sprintf("service:%s", service)}, 1)
	}
	ls.serviceCount = make(map[string]int64)
	log.WithField("total_spans", totalCount).Debug("Checkpointing flushed spans for Lightstep")
}
<|endoftext|>


<commit_before>package mysqlwarmer

import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"
	"testing"

	"github.com/lestrrat/go-test-mysqld"
	"github.com/stretchr/testify/assert"
)

var (
	testMysqld *mysqltest.TestMysqld
	opts       = Options{
		DataBase: "test",
		Host:     "localhost",
		User:     "root",
		Password: "",
		Port:     13306,
	}
)

func TestMain(m *testing.M) {
	os.Exit(runTests(m))
}

func createTestDataBase() {
	db, err := sql.Open("mysql", fmt.Sprintf("root:@tcp(%s:%d)/", opts.Host, opts.Port))
	if err != nil {
		log.Fatalln(err)
	}
	defer db.Close()

	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS  `test`"); err != nil {
		log.Fatalln(err)
	}

	if _, err := db.Exec("USE `test`"); err != nil {
		log.Fatalln(err)
	}
}

func runTests(m *testing.M) int {
	config := mysqltest.NewConfig()
	config.SkipNetworking = false
	config.Port = opts.Port
	config.BindAddress = opts.Host
	mysqld, err := mysqltest.NewMysqld(config)
	if err != nil {
		log.Fatalln(err)
	}
	defer mysqld.Stop()

	testMysqld = mysqld

	createTestDataBase()

	return m.Run()
}

func TestGetDSN(t *testing.T) {
	assert.Equal(
		t,
		opts.getDSN(),
		fmt.Sprintf("%s:%s@(%s:%d)/%s", opts.User, opts.Password, opts.Host, opts.Port, opts.DataBase),
		"should be equal",
	)
}

func createTestTable() {
	db, err := sql.Open("mysql", testMysqld.DSN())
	if err != nil {
		log.Fatalln(err)
	}
	defer db.Close()

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS `myisam_table` (`id` INTEGER, INDEX(`id`)) ENGINE=MyISAM"); err != nil {
		log.Fatalln(err)
	}

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS `innodb_table` (`id` INTEGER) ENGINE=InnoDB"); err != nil {
		log.Fatalln(err)
	}
}

func TestGetTables(t *testing.T) {
	createTestTable()

	var myisamTables []string
	myisamTables = append(myisamTables, "myisam_table")
	tables, err := opts.getTables("myisam")
	if err != nil {
		log.Fatalln(err)
	}

	assert.Equal(
		t,
		tables,
		myisamTables,
		"should be equal",
	)
}

func checkVersion() bool {
	out, err := exec.Command("mysql", "--version").Output()
	if err != nil {
		log.Fatalln(err)
	}

	r := regexp.MustCompile("5.7")
	return r.MatchString(fmt.Sprintf("%s", out))
}

func TestPreload(t *testing.T) {
	myisamtables := []string{"myisam_table"}

	db, err := sql.Open("mysql", testMysqld.DSN())
	if err != nil {
		log.Fatalln(err)
	}
	defer db.Close()

	if checkVersion() {
		if _, err = db.Exec("SET GLOBAL show_compatibility_56=ON"); err != nil {
			log.Fatalln(err)
		}
	}

	var keysCount int
	err = db.QueryRow("SELECT `VARIABLE_VALUE` FROM information_schema.GLOBAL_STATUS  WHERE VARIABLE_NAME = 'com_preload_keys'").Scan(&keysCount)
	if err != nil {
		log.Fatalln(err)
	}

	assert.Equal(
		t,
		keysCount,
		0,
		"has not cache",
	)

	if err = opts.preload(myisamtables); err != nil {
		log.Fatalln(err)
	}

	err = db.QueryRow("SELECT `VARIABLE_VALUE` FROM information_schema.GLOBAL_STATUS  WHERE VARIABLE_NAME = 'com_preload_keys'").Scan(&keysCount)
	if err != nil {
		log.Fatalln(err)
	}

	assert.Equal(
		t,
		keysCount,
		1,
		"has cache",
	)
}

func TestRun(t *testing.T) {
	args := []string{
		"--host",
		opts.Host,
		"--user",
		opts.User,
		"--port",
		fmt.Sprintf("%d", opts.Port),
		"--database",
		opts.DataBase,
	}
	assert.Nil(t, Run(args), "success")
}
<commit_msg>:bug: fix regex bug<commit_after>package mysqlwarmer

import (
	"bytes"
	"database/sql"
	"fmt"
	"log"
	"os"
	"os/exec"
	"testing"

	"github.com/lestrrat/go-test-mysqld"
	"github.com/stretchr/testify/assert"
)

var (
	testMysqld *mysqltest.TestMysqld
	opts       = Options{
		DataBase: "test",
		Host:     "localhost",
		User:     "root",
		Password: "",
		Port:     13306,
	}
)

func TestMain(m *testing.M) {
	os.Exit(runTests(m))
}

func createTestDataBase() {
	db, err := sql.Open("mysql", fmt.Sprintf("root:@tcp(%s:%d)/", opts.Host, opts.Port))
	if err != nil {
		log.Fatalln(err)
	}
	defer db.Close()

	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS  `test`"); err != nil {
		log.Fatalln(err)
	}

	if _, err := db.Exec("USE `test`"); err != nil {
		log.Fatalln(err)
	}
}

func runTests(m *testing.M) int {
	config := mysqltest.NewConfig()
	config.SkipNetworking = false
	config.Port = opts.Port
	config.BindAddress = opts.Host
	mysqld, err := mysqltest.NewMysqld(config)
	if err != nil {
		log.Fatalln(err)
	}
	defer mysqld.Stop()

	testMysqld = mysqld

	createTestDataBase()

	return m.Run()
}

func TestGetDSN(t *testing.T) {
	assert.Equal(
		t,
		opts.getDSN(),
		fmt.Sprintf("%s:%s@(%s:%d)/%s", opts.User, opts.Password, opts.Host, opts.Port, opts.DataBase),
		"should be equal",
	)
}

func createTestTable() {
	db, err := sql.Open("mysql", testMysqld.DSN())
	if err != nil {
		log.Fatalln(err)
	}
	defer db.Close()

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS `myisam_table` (`id` INTEGER, INDEX(`id`)) ENGINE=MyISAM"); err != nil {
		log.Fatalln(err)
	}

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS `innodb_table` (`id` INTEGER) ENGINE=InnoDB"); err != nil {
		log.Fatalln(err)
	}
}

func TestGetTables(t *testing.T) {
	createTestTable()

	var myisamTables []string
	myisamTables = append(myisamTables, "myisam_table")
	tables, err := opts.getTables("myisam")
	if err != nil {
		log.Fatalln(err)
	}

	assert.Equal(
		t,
		tables,
		myisamTables,
		"should be equal",
	)
}

func checkVersion() bool {
	out, err := exec.Command("mysql", "--version").Output()
	if err != nil {
		log.Fatalln(err)
	}

	return bytes.Contains(out, []byte("5.7"))
}

func TestPreload(t *testing.T) {
	myisamtables := []string{"myisam_table"}

	db, err := sql.Open("mysql", testMysqld.DSN())
	if err != nil {
		log.Fatalln(err)
	}
	defer db.Close()

	if checkVersion() {
		if _, err = db.Exec("SET GLOBAL show_compatibility_56=ON"); err != nil {
			log.Fatalln(err)
		}
	}

	var keysCount int
	err = db.QueryRow("SELECT `VARIABLE_VALUE` FROM information_schema.GLOBAL_STATUS  WHERE VARIABLE_NAME = 'com_preload_keys'").Scan(&keysCount)
	if err != nil {
		log.Fatalln(err)
	}

	assert.Equal(
		t,
		keysCount,
		0,
		"has not cache",
	)

	if err = opts.preload(myisamtables); err != nil {
		log.Fatalln(err)
	}

	err = db.QueryRow("SELECT `VARIABLE_VALUE` FROM information_schema.GLOBAL_STATUS  WHERE VARIABLE_NAME = 'com_preload_keys'").Scan(&keysCount)
	if err != nil {
		log.Fatalln(err)
	}

	assert.Equal(
		t,
		keysCount,
		1,
		"has cache",
	)
}

func TestRun(t *testing.T) {
	args := []string{
		"--host",
		opts.Host,
		"--user",
		opts.User,
		"--port",
		fmt.Sprintf("%d", opts.Port),
		"--database",
		opts.DataBase,
	}
	assert.Nil(t, Run(args), "success")
}
<|endoftext|>