<commit_before>package untar

import (
	"archive/tar"
	"errors"
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	"io"
	"io/ioutil"
	"path"
	"syscall"
)

type decoderData struct {
	nextInodeNumber   uint64
	fileSystem        filesystem.FileSystem
	regularInodeTable map[string]uint64
	symlinkInodeTable map[string]uint64
	inodeTable        map[string]uint64
	directoryTable    map[string]*filesystem.Directory
}

func decode(tarReader *tar.Reader, dataHandler DataHandler,
	filter *filter.Filter) (*filesystem.FileSystem, error) {
	var decoderData decoderData
	decoderData.regularInodeTable = make(map[string]uint64)
	decoderData.symlinkInodeTable = make(map[string]uint64)
	decoderData.inodeTable = make(map[string]uint64)
	decoderData.directoryTable = make(map[string]*filesystem.Directory)
	fileSystem := &decoderData.fileSystem
	fileSystem.RegularInodeTable = make(filesystem.RegularInodeTable)
	fileSystem.SymlinkInodeTable = make(filesystem.SymlinkInodeTable)
	fileSystem.InodeTable = make(filesystem.InodeTable)
	// Create a container directory for the top-level directory.
	var containerDir filesystem.Directory
	decoderData.directoryTable["/"] = &containerDir
	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		header.Name = normaliseFilename(header.Name)
		if filter.Match(header.Name) {
			continue
		}
		err = decoderData.addHeader(tarReader, dataHandler, header)
		if err != nil {
			return nil, err
		}
	}
	fileSystem.Directory = *containerDir.DirectoryList[0]
	sortDirectory(&fileSystem.Directory)
	fileSystem.DirectoryCount = uint64(len(decoderData.directoryTable))
	fileSystem.RebuildInodePointers()
	fileSystem.ComputeTotalDataBytes()
	return fileSystem, nil
}

func normaliseFilename(filename string) string {
	if filename[:2] == "./" {
		filename = filename[1:]
	} else if filename[0] != '/' {
		filename = "/" + filename
	}
	length := len(filename)
	if length > 1 && filename[length-1] == '/' {
		filename = filename[:length-1]
	}
	return filename
}

func (decoderData *decoderData) addHeader(tarReader *tar.Reader,
	dataHandler DataHandler, header *tar.Header) error {
	parentDir, ok := decoderData.directoryTable[path.Dir(header.Name)]
	if !ok {
		return errors.New(fmt.Sprintf(
			"No parent directory found for: %s", header.Name))
	}
	leafName := path.Base(header.Name)
	if header.Typeflag == tar.TypeReg || header.Typeflag == tar.TypeRegA {
		return decoderData.addRegularFile(tarReader, dataHandler, header,
			parentDir, leafName)
	} else if header.Typeflag == tar.TypeLink {
		return decoderData.addDirectory(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeSymlink {
		return decoderData.addSymlink(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeChar {
		return decoderData.addFile(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeBlock {
		return decoderData.addFile(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeDir {
		return decoderData.addDirectory(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeFifo {
		return decoderData.addFile(header, parentDir, leafName)
	} else {
		return errors.New(fmt.Sprintf("Unsupported file type: %v",
			header.Typeflag))
	}
	return nil
}

func (decoderData *decoderData) addRegularFile(tarReader *tar.Reader,
	dataHandler DataHandler, header *tar.Header, parent *filesystem.Directory,
	name string) error {
	var newInode filesystem.RegularInode
	newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
		syscall.S_IFREG)
	newInode.Uid = uint32(header.Uid)
	newInode.Gid = uint32(header.Gid)
	newInode.MtimeNanoSeconds = int32(header.ModTime.Nanosecond())
	newInode.MtimeSeconds = header.ModTime.Unix()
	newInode.Size = uint64(header.Size)
	if header.Size > 0 {
		data, err := ioutil.ReadAll(tarReader)
		if err != nil {
			return errors.New("error reading file data" + err.Error())
		}
		if int64(len(data)) != header.Size {
			return errors.New(fmt.Sprintf(
				"failed to read file data, wanted: %d, got: %d bytes",
				header.Size, len(data)))
		}
		newInode.Hash, err = dataHandler.HandleData(data)
		if err != nil {
			return err
		}
	}
	decoderData.regularInodeTable[header.Name] = decoderData.nextInodeNumber
	decoderData.fileSystem.RegularInodeTable[decoderData.nextInodeNumber] =
		&newInode
	var newEntry filesystem.RegularFile
	newEntry.Name = name
	newEntry.InodeNumber = decoderData.nextInodeNumber
	parent.RegularFileList = append(parent.RegularFileList, &newEntry)
	decoderData.nextInodeNumber++
	return nil
}

func (decoderData *decoderData) addDirectory(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	var newEntry filesystem.Directory
	newEntry.Name = name
	newEntry.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
		syscall.S_IFDIR)
	newEntry.Uid = uint32(header.Uid)
	newEntry.Gid = uint32(header.Gid)
	parent.DirectoryList = append(parent.DirectoryList, &newEntry)
	decoderData.directoryTable[header.Name] = &newEntry
	return nil
}

func (decoderData *decoderData) addHardlink(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	header.Linkname = normaliseFilename(header.Linkname)
	if inum, ok := decoderData.regularInodeTable[header.Linkname]; ok {
		var newEntry filesystem.RegularFile
		newEntry.Name = name
		newEntry.InodeNumber = inum
		parent.RegularFileList = append(parent.RegularFileList, &newEntry)
	} else if inum, ok := decoderData.symlinkInodeTable[header.Linkname]; ok {
		var newEntry filesystem.Symlink
		newEntry.Name = name
		newEntry.InodeNumber = inum
		parent.SymlinkList = append(parent.SymlinkList, &newEntry)
	} else if inum, ok := decoderData.inodeTable[header.Linkname]; ok {
		var newEntry filesystem.File
		newEntry.Name = name
		newEntry.InodeNumber = inum
		parent.FileList = append(parent.FileList, &newEntry)
	} else {
		return errors.New(fmt.Sprintf("missing hardlink target: %s",
			header.Linkname))
	}
	return nil
}

func (decoderData *decoderData) addSymlink(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	var newInode filesystem.SymlinkInode
	newInode.Uid = uint32(header.Uid)
	newInode.Gid = uint32(header.Gid)
	newInode.Symlink = header.Linkname
	decoderData.symlinkInodeTable[header.Name] = decoderData.nextInodeNumber
	decoderData.fileSystem.SymlinkInodeTable[decoderData.nextInodeNumber] =
		&newInode
	var newEntry filesystem.Symlink
	newEntry.Name = name
	newEntry.InodeNumber = decoderData.nextInodeNumber
	parent.SymlinkList = append(parent.SymlinkList, &newEntry)
	decoderData.nextInodeNumber++
	return nil
}

func (decoderData *decoderData) addFile(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	var newInode filesystem.Inode
	if header.Typeflag == tar.TypeChar {
		newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
			syscall.S_IFCHR)
	} else if header.Typeflag == tar.TypeBlock {
		newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
			syscall.S_IFBLK)
	} else if header.Typeflag == tar.TypeFifo {
		newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
			syscall.S_IFIFO)
	} else {
		return errors.New(fmt.Sprintf("unsupported type: %v", header.Typeflag))
	}
	newInode.Uid = uint32(header.Uid)
	newInode.Gid = uint32(header.Gid)
	newInode.MtimeNanoSeconds = int32(header.ModTime.Nanosecond())
	newInode.MtimeSeconds = header.ModTime.Unix()
	if header.Devminor > 255 {
		return errors.New(fmt.Sprintf("minor device number: %d too large",
			header.Devminor))
	}
	newInode.Rdev = uint64(header.Devmajor>>8 | header.Devminor)
	decoderData.inodeTable[header.Name] = decoderData.nextInodeNumber
	decoderData.fileSystem.InodeTable[decoderData.nextInodeNumber] =
		&newInode
	var newEntry filesystem.File
	newEntry.Name = name
	newEntry.InodeNumber = decoderData.nextInodeNumber
	parent.FileList = append(parent.FileList, &newEntry)
	decoderData.nextInodeNumber++
	return nil
}
<commit_msg>Fix bug in tar decoder: device numbers were mangled.<commit_after>package untar

import (
	"archive/tar"
	"errors"
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	"io"
	"io/ioutil"
	"path"
	"syscall"
)

type decoderData struct {
	nextInodeNumber   uint64
	fileSystem        filesystem.FileSystem
	regularInodeTable map[string]uint64
	symlinkInodeTable map[string]uint64
	inodeTable        map[string]uint64
	directoryTable    map[string]*filesystem.Directory
}

func decode(tarReader *tar.Reader, dataHandler DataHandler,
	filter *filter.Filter) (*filesystem.FileSystem, error) {
	var decoderData decoderData
	decoderData.regularInodeTable = make(map[string]uint64)
	decoderData.symlinkInodeTable = make(map[string]uint64)
	decoderData.inodeTable = make(map[string]uint64)
	decoderData.directoryTable = make(map[string]*filesystem.Directory)
	fileSystem := &decoderData.fileSystem
	fileSystem.RegularInodeTable = make(filesystem.RegularInodeTable)
	fileSystem.SymlinkInodeTable = make(filesystem.SymlinkInodeTable)
	fileSystem.InodeTable = make(filesystem.InodeTable)
	// Create a container directory for the top-level directory.
	var containerDir filesystem.Directory
	decoderData.directoryTable["/"] = &containerDir
	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		header.Name = normaliseFilename(header.Name)
		if filter.Match(header.Name) {
			continue
		}
		err = decoderData.addHeader(tarReader, dataHandler, header)
		if err != nil {
			return nil, err
		}
	}
	fileSystem.Directory = *containerDir.DirectoryList[0]
	sortDirectory(&fileSystem.Directory)
	fileSystem.DirectoryCount = uint64(len(decoderData.directoryTable))
	fileSystem.RebuildInodePointers()
	fileSystem.ComputeTotalDataBytes()
	return fileSystem, nil
}

func normaliseFilename(filename string) string {
	if filename[:2] == "./" {
		filename = filename[1:]
	} else if filename[0] != '/' {
		filename = "/" + filename
	}
	length := len(filename)
	if length > 1 && filename[length-1] == '/' {
		filename = filename[:length-1]
	}
	return filename
}

func (decoderData *decoderData) addHeader(tarReader *tar.Reader,
	dataHandler DataHandler, header *tar.Header) error {
	parentDir, ok := decoderData.directoryTable[path.Dir(header.Name)]
	if !ok {
		return errors.New(fmt.Sprintf(
			"No parent directory found for: %s", header.Name))
	}
	leafName := path.Base(header.Name)
	if header.Typeflag == tar.TypeReg || header.Typeflag == tar.TypeRegA {
		return decoderData.addRegularFile(tarReader, dataHandler, header,
			parentDir, leafName)
	} else if header.Typeflag == tar.TypeLink {
		return decoderData.addDirectory(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeSymlink {
		return decoderData.addSymlink(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeChar {
		return decoderData.addFile(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeBlock {
		return decoderData.addFile(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeDir {
		return decoderData.addDirectory(header, parentDir, leafName)
	} else if header.Typeflag == tar.TypeFifo {
		return decoderData.addFile(header, parentDir, leafName)
	} else {
		return errors.New(fmt.Sprintf("Unsupported file type: %v",
			header.Typeflag))
	}
	return nil
}

func (decoderData *decoderData) addRegularFile(tarReader *tar.Reader,
	dataHandler DataHandler, header *tar.Header, parent *filesystem.Directory,
	name string) error {
	var newInode filesystem.RegularInode
	newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
		syscall.S_IFREG)
	newInode.Uid = uint32(header.Uid)
	newInode.Gid = uint32(header.Gid)
	newInode.MtimeNanoSeconds = int32(header.ModTime.Nanosecond())
	newInode.MtimeSeconds = header.ModTime.Unix()
	newInode.Size = uint64(header.Size)
	if header.Size > 0 {
		data, err := ioutil.ReadAll(tarReader)
		if err != nil {
			return errors.New("error reading file data" + err.Error())
		}
		if int64(len(data)) != header.Size {
			return errors.New(fmt.Sprintf(
				"failed to read file data, wanted: %d, got: %d bytes",
				header.Size, len(data)))
		}
		newInode.Hash, err = dataHandler.HandleData(data)
		if err != nil {
			return err
		}
	}
	decoderData.regularInodeTable[header.Name] = decoderData.nextInodeNumber
	decoderData.fileSystem.RegularInodeTable[decoderData.nextInodeNumber] =
		&newInode
	var newEntry filesystem.RegularFile
	newEntry.Name = name
	newEntry.InodeNumber = decoderData.nextInodeNumber
	parent.RegularFileList = append(parent.RegularFileList, &newEntry)
	decoderData.nextInodeNumber++
	return nil
}

func (decoderData *decoderData) addDirectory(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	var newEntry filesystem.Directory
	newEntry.Name = name
	newEntry.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
		syscall.S_IFDIR)
	newEntry.Uid = uint32(header.Uid)
	newEntry.Gid = uint32(header.Gid)
	parent.DirectoryList = append(parent.DirectoryList, &newEntry)
	decoderData.directoryTable[header.Name] = &newEntry
	return nil
}

func (decoderData *decoderData) addHardlink(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	header.Linkname = normaliseFilename(header.Linkname)
	if inum, ok := decoderData.regularInodeTable[header.Linkname]; ok {
		var newEntry filesystem.RegularFile
		newEntry.Name = name
		newEntry.InodeNumber = inum
		parent.RegularFileList = append(parent.RegularFileList, &newEntry)
	} else if inum, ok := decoderData.symlinkInodeTable[header.Linkname]; ok {
		var newEntry filesystem.Symlink
		newEntry.Name = name
		newEntry.InodeNumber = inum
		parent.SymlinkList = append(parent.SymlinkList, &newEntry)
	} else if inum, ok := decoderData.inodeTable[header.Linkname]; ok {
		var newEntry filesystem.File
		newEntry.Name = name
		newEntry.InodeNumber = inum
		parent.FileList = append(parent.FileList, &newEntry)
	} else {
		return errors.New(fmt.Sprintf("missing hardlink target: %s",
			header.Linkname))
	}
	return nil
}

func (decoderData *decoderData) addSymlink(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	var newInode filesystem.SymlinkInode
	newInode.Uid = uint32(header.Uid)
	newInode.Gid = uint32(header.Gid)
	newInode.Symlink = header.Linkname
	decoderData.symlinkInodeTable[header.Name] = decoderData.nextInodeNumber
	decoderData.fileSystem.SymlinkInodeTable[decoderData.nextInodeNumber] =
		&newInode
	var newEntry filesystem.Symlink
	newEntry.Name = name
	newEntry.InodeNumber = decoderData.nextInodeNumber
	parent.SymlinkList = append(parent.SymlinkList, &newEntry)
	decoderData.nextInodeNumber++
	return nil
}

func (decoderData *decoderData) addFile(header *tar.Header,
	parent *filesystem.Directory, name string) error {
	var newInode filesystem.Inode
	if header.Typeflag == tar.TypeChar {
		newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
			syscall.S_IFCHR)
	} else if header.Typeflag == tar.TypeBlock {
		newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
			syscall.S_IFBLK)
	} else if header.Typeflag == tar.TypeFifo {
		newInode.Mode = filesystem.FileMode((header.Mode & ^syscall.S_IFMT) |
			syscall.S_IFIFO)
	} else {
		return errors.New(fmt.Sprintf("unsupported type: %v", header.Typeflag))
	}
	newInode.Uid = uint32(header.Uid)
	newInode.Gid = uint32(header.Gid)
	newInode.MtimeNanoSeconds = int32(header.ModTime.Nanosecond())
	newInode.MtimeSeconds = header.ModTime.Unix()
	if header.Devminor > 255 {
		return errors.New(fmt.Sprintf("minor device number: %d too large",
			header.Devminor))
	}
	newInode.Rdev = uint64(header.Devmajor<<8 | header.Devminor)
	decoderData.inodeTable[header.Name] = decoderData.nextInodeNumber
	decoderData.fileSystem.InodeTable[decoderData.nextInodeNumber] =
		&newInode
	var newEntry filesystem.File
	newEntry.Name = name
	newEntry.InodeNumber = decoderData.nextInodeNumber
	parent.FileList = append(parent.FileList, &newEntry)
	decoderData.nextInodeNumber++
	return nil
}
<|endoftext|>


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

package main

import (
	"github.com/globocom/config"
	"github.com/globocom/tsuru/cmd"
	"github.com/globocom/tsuru/provision"
	_ "github.com/globocom/tsuru/provision/docker"
	_ "github.com/globocom/tsuru/provision/juju"
	"os"
)

const defaultConfigPath = "/etc/tsuru/tsuru.conf"

func buildManager() *cmd.Manager {
	m := cmd.NewManager("tsr", "0.2.9", "", os.Stdout, os.Stderr, os.Stdin)
	m.Register(&tsrCommand{Command: &apiCmd{}})
	m.Register(&tsrCommand{Command: &collectorCmd{}})
	m.Register(&tsrCommand{Command: tokenCmd{}})
	m.Register(&tsrCommand{Command: &healerCmd{}})
	registerProvisionersCommands(m)
	return m
}

func registerProvisionersCommands(m *cmd.Manager) {
	provisioners := provision.Registry()
	for _, p := range provisioners {
		if c, ok := p.(provision.Commandable); ok {
			commands := c.Commands()
			for _, cmd := range commands {
				m.Register(&tsrCommand{Command: cmd})
			}
		}
	}
}

func main() {
	config.ReadConfigFile(defaultConfigPath)
	m := buildManager()
	m.Run(os.Args[1:])
}
<commit_msg>bump to 0.2.10.<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"github.com/globocom/config"
	"github.com/globocom/tsuru/cmd"
	"github.com/globocom/tsuru/provision"
	_ "github.com/globocom/tsuru/provision/docker"
	_ "github.com/globocom/tsuru/provision/juju"
	"os"
)

const defaultConfigPath = "/etc/tsuru/tsuru.conf"

func buildManager() *cmd.Manager {
	m := cmd.NewManager("tsr", "0.2.10", "", os.Stdout, os.Stderr, os.Stdin)
	m.Register(&tsrCommand{Command: &apiCmd{}})
	m.Register(&tsrCommand{Command: &collectorCmd{}})
	m.Register(&tsrCommand{Command: tokenCmd{}})
	m.Register(&tsrCommand{Command: &healerCmd{}})
	registerProvisionersCommands(m)
	return m
}

func registerProvisionersCommands(m *cmd.Manager) {
	provisioners := provision.Registry()
	for _, p := range provisioners {
		if c, ok := p.(provision.Commandable); ok {
			commands := c.Commands()
			for _, cmd := range commands {
				m.Register(&tsrCommand{Command: cmd})
			}
		}
	}
}

func main() {
	config.ReadConfigFile(defaultConfigPath)
	m := buildManager()
	m.Run(os.Args[1:])
}
<|endoftext|>


<commit_before>package boilingcore

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"sort"
	"strings"
	"text/template"

	"github.com/pkg/errors"
	"github.com/volatiletech/sqlboiler/drivers"
	"github.com/volatiletech/sqlboiler/strmangle"
)

// templateData for sqlboiler templates
type templateData struct {
	Tables []drivers.Table
	Table  drivers.Table

	// Controls what names are output
	PkgName string
	Schema  string

	// Controls which code is output (mysql vs pqsl ...)
	DriverName string
	Dialect    drivers.Dialect

	// LQ and RQ contain a quoted quote that allows us to write
	// the templates easier.
	LQ string
	RQ string

	// Turn off auto timestamps or hook generation
	NoHooks          bool
	NoAutoTimestamps bool

	// Tags control which
	Tags []string

	// Generate struct tags as camelCase or snake_case
	StructTagCasing string

	// StringFuncs are usable in templates with stringMap
	StringFuncs map[string]func(string) string
}

func (t templateData) Quotes(s string) string {
	return fmt.Sprintf("%s%s%s", t.LQ, s, t.RQ)
}

func (t templateData) SchemaTable(table string) string {
	return strmangle.SchemaTable(t.LQ, t.RQ, t.DriverName, t.Schema, table)
}

type templateList struct {
	*template.Template
}

type templateNameList []string

func (t templateNameList) Len() int {
	return len(t)
}

func (t templateNameList) Swap(k, j int) {
	t[k], t[j] = t[j], t[k]
}

func (t templateNameList) Less(k, j int) bool {
	// Make sure "struct" goes to the front
	if t[k] == "struct.tpl" {
		return true
	}

	res := strings.Compare(t[k], t[j])
	if res <= 0 {
		return true
	}

	return false
}

// Templates returns the name of all the templates defined in the template list
func (t templateList) Templates() []string {
	tplList := t.Template.Templates()

	if len(tplList) == 0 {
		return nil
	}

	ret := make([]string, 0, len(tplList))
	for _, tpl := range tplList {
		if name := tpl.Name(); strings.HasSuffix(name, ".tpl") {
			ret = append(ret, name)
		}
	}

	sort.Sort(templateNameList(ret))

	return ret
}

// loadTemplates loads all of the template files in the specified directory.
func loadTemplates(dir string) (*templateList, error) {
	pattern := filepath.Join(dir, "*.tpl")
	tpl, err := template.New("").Funcs(templateFunctions).ParseGlob(pattern)

	if err != nil {
		return nil, err
	}

	return &templateList{Template: tpl}, err
}

// loadTemplate loads a single template file
func loadTemplate(dir string, filename string) (*template.Template, error) {
	pattern := filepath.Join(dir, filename)
	tpl, err := template.New("").Funcs(templateFunctions).ParseFiles(pattern)

	if err != nil {
		return nil, err
	}

	return tpl.Lookup(filename), err
}

// replaceTemplate finds the template matching with name and replaces its
// contents with the contents of the template located at filename
func replaceTemplate(tpl *template.Template, name, filename string) error {
	if tpl == nil {
		return fmt.Errorf("template for %s is nil", name)
	}

	b, err := ioutil.ReadFile(filename)
	if err != nil {
		return errors.Wrapf(err, "failed reading template file: %s", filename)
	}

	if tpl, err = tpl.New(name).Funcs(templateFunctions).Parse(string(b)); err != nil {
		return errors.Wrapf(err, "failed to parse template file: %s", filename)
	}

	return nil
}

// set is to stop duplication from named enums, allowing a template loop
// to keep some state
type once map[string]struct{}

func newOnce() once {
	return make(once)
}

func (o once) Has(s string) bool {
	_, ok := o[s]
	return ok
}

func (o once) Put(s string) bool {
	if _, ok := o[s]; ok {
		return false
	}

	o[s] = struct{}{}
	return true
}

// templateStringMappers are placed into the data to make it easy to use the
// stringMap function.
var templateStringMappers = map[string]func(string) string{
	// String ops
	"quoteWrap":       func(a string) string { return fmt.Sprintf(`"%s"`, a) },
	"replaceReserved": strmangle.ReplaceReservedWords,

	// Casing
	"titleCase": strmangle.TitleCase,
	"camelCase": strmangle.CamelCase,
}

// templateFunctions is a map of all the functions that get passed into the
// templates. If you wish to pass a new function into your own template,
// add a function pointer here.
var templateFunctions = template.FuncMap{
	// String ops
	"quoteWrap": func(s string) string { return fmt.Sprintf(`"%s"`, s) },
	"id":        strmangle.Identifier,

	// Pluralization
	"singular": strmangle.Singular,
	"plural":   strmangle.Plural,

	// Casing
	"titleCase": strmangle.TitleCase,
	"camelCase": strmangle.CamelCase,

	// String Slice ops
	"join":               func(sep string, slice []string) string { return strings.Join(slice, sep) },
	"joinSlices":         strmangle.JoinSlices,
	"stringMap":          strmangle.StringMap,
	"prefixStringSlice":  strmangle.PrefixStringSlice,
	"containsAny":        strmangle.ContainsAny,
	"generateTags":       strmangle.GenerateTags,
	"generateIgnoreTags": strmangle.GenerateIgnoreTags,

	// Enum ops
	"parseEnumName":       strmangle.ParseEnumName,
	"parseEnumVals":       strmangle.ParseEnumVals,
	"isEnumNormal":        strmangle.IsEnumNormal,
	"shouldTitleCaseEnum": strmangle.ShouldTitleCaseEnum,
	"onceNew":             newOnce,
	"oncePut":             once.Put,
	"onceHas":             once.Has,

	// String Map ops
	"makeStringMap": strmangle.MakeStringMap,

	// Set operations
	"setInclude": strmangle.SetInclude,

	// Database related mangling
	"whereClause": strmangle.WhereClause,

	// Relationship text helpers
	"txtsFromFKey":     txtsFromFKey,
	"txtsFromOneToOne": txtsFromOneToOne,
	"txtsFromToMany":   txtsFromToMany,

	// dbdrivers ops
	"filterColumnsByAuto":    drivers.FilterColumnsByAuto,
	"filterColumnsByDefault": drivers.FilterColumnsByDefault,
	"filterColumnsByEnum":    drivers.FilterColumnsByEnum,
	"sqlColDefinitions":      drivers.SQLColDefinitions,
	"columnNames":            drivers.ColumnNames,
	"columnDBTypes":          drivers.ColumnDBTypes,
	"getTable":               drivers.GetTable,
}
<commit_msg>Fix comment<commit_after>package boilingcore

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"sort"
	"strings"
	"text/template"

	"github.com/pkg/errors"
	"github.com/volatiletech/sqlboiler/drivers"
	"github.com/volatiletech/sqlboiler/strmangle"
)

// templateData for sqlboiler templates
type templateData struct {
	Tables []drivers.Table
	Table  drivers.Table

	// Controls what names are output
	PkgName string
	Schema  string

	// Controls which code is output (mysql vs pqsl ...)
	DriverName string
	Dialect    drivers.Dialect

	// LQ and RQ contain a quoted quote that allows us to write
	// the templates more easily.
	LQ string
	RQ string

	// Turn off auto timestamps or hook generation
	NoHooks          bool
	NoAutoTimestamps bool

	// Tags control which
	Tags []string

	// Generate struct tags as camelCase or snake_case
	StructTagCasing string

	// StringFuncs are usable in templates with stringMap
	StringFuncs map[string]func(string) string
}

func (t templateData) Quotes(s string) string {
	return fmt.Sprintf("%s%s%s", t.LQ, s, t.RQ)
}

func (t templateData) SchemaTable(table string) string {
	return strmangle.SchemaTable(t.LQ, t.RQ, t.DriverName, t.Schema, table)
}

type templateList struct {
	*template.Template
}

type templateNameList []string

func (t templateNameList) Len() int {
	return len(t)
}

func (t templateNameList) Swap(k, j int) {
	t[k], t[j] = t[j], t[k]
}

func (t templateNameList) Less(k, j int) bool {
	// Make sure "struct" goes to the front
	if t[k] == "struct.tpl" {
		return true
	}

	res := strings.Compare(t[k], t[j])
	if res <= 0 {
		return true
	}

	return false
}

// Templates returns the name of all the templates defined in the template list
func (t templateList) Templates() []string {
	tplList := t.Template.Templates()

	if len(tplList) == 0 {
		return nil
	}

	ret := make([]string, 0, len(tplList))
	for _, tpl := range tplList {
		if name := tpl.Name(); strings.HasSuffix(name, ".tpl") {
			ret = append(ret, name)
		}
	}

	sort.Sort(templateNameList(ret))

	return ret
}

// loadTemplates loads all of the template files in the specified directory.
func loadTemplates(dir string) (*templateList, error) {
	pattern := filepath.Join(dir, "*.tpl")
	tpl, err := template.New("").Funcs(templateFunctions).ParseGlob(pattern)

	if err != nil {
		return nil, err
	}

	return &templateList{Template: tpl}, err
}

// loadTemplate loads a single template file
func loadTemplate(dir string, filename string) (*template.Template, error) {
	pattern := filepath.Join(dir, filename)
	tpl, err := template.New("").Funcs(templateFunctions).ParseFiles(pattern)

	if err != nil {
		return nil, err
	}

	return tpl.Lookup(filename), err
}

// replaceTemplate finds the template matching with name and replaces its
// contents with the contents of the template located at filename
func replaceTemplate(tpl *template.Template, name, filename string) error {
	if tpl == nil {
		return fmt.Errorf("template for %s is nil", name)
	}

	b, err := ioutil.ReadFile(filename)
	if err != nil {
		return errors.Wrapf(err, "failed reading template file: %s", filename)
	}

	if tpl, err = tpl.New(name).Funcs(templateFunctions).Parse(string(b)); err != nil {
		return errors.Wrapf(err, "failed to parse template file: %s", filename)
	}

	return nil
}

// set is to stop duplication from named enums, allowing a template loop
// to keep some state
type once map[string]struct{}

func newOnce() once {
	return make(once)
}

func (o once) Has(s string) bool {
	_, ok := o[s]
	return ok
}

func (o once) Put(s string) bool {
	if _, ok := o[s]; ok {
		return false
	}

	o[s] = struct{}{}
	return true
}

// templateStringMappers are placed into the data to make it easy to use the
// stringMap function.
var templateStringMappers = map[string]func(string) string{
	// String ops
	"quoteWrap":       func(a string) string { return fmt.Sprintf(`"%s"`, a) },
	"replaceReserved": strmangle.ReplaceReservedWords,

	// Casing
	"titleCase": strmangle.TitleCase,
	"camelCase": strmangle.CamelCase,
}

// templateFunctions is a map of all the functions that get passed into the
// templates. If you wish to pass a new function into your own template,
// add a function pointer here.
var templateFunctions = template.FuncMap{
	// String ops
	"quoteWrap": func(s string) string { return fmt.Sprintf(`"%s"`, s) },
	"id":        strmangle.Identifier,

	// Pluralization
	"singular": strmangle.Singular,
	"plural":   strmangle.Plural,

	// Casing
	"titleCase": strmangle.TitleCase,
	"camelCase": strmangle.CamelCase,

	// String Slice ops
	"join":               func(sep string, slice []string) string { return strings.Join(slice, sep) },
	"joinSlices":         strmangle.JoinSlices,
	"stringMap":          strmangle.StringMap,
	"prefixStringSlice":  strmangle.PrefixStringSlice,
	"containsAny":        strmangle.ContainsAny,
	"generateTags":       strmangle.GenerateTags,
	"generateIgnoreTags": strmangle.GenerateIgnoreTags,

	// Enum ops
	"parseEnumName":       strmangle.ParseEnumName,
	"parseEnumVals":       strmangle.ParseEnumVals,
	"isEnumNormal":        strmangle.IsEnumNormal,
	"shouldTitleCaseEnum": strmangle.ShouldTitleCaseEnum,
	"onceNew":             newOnce,
	"oncePut":             once.Put,
	"onceHas":             once.Has,

	// String Map ops
	"makeStringMap": strmangle.MakeStringMap,

	// Set operations
	"setInclude": strmangle.SetInclude,

	// Database related mangling
	"whereClause": strmangle.WhereClause,

	// Relationship text helpers
	"txtsFromFKey":     txtsFromFKey,
	"txtsFromOneToOne": txtsFromOneToOne,
	"txtsFromToMany":   txtsFromToMany,

	// dbdrivers ops
	"filterColumnsByAuto":    drivers.FilterColumnsByAuto,
	"filterColumnsByDefault": drivers.FilterColumnsByDefault,
	"filterColumnsByEnum":    drivers.FilterColumnsByEnum,
	"sqlColDefinitions":      drivers.SQLColDefinitions,
	"columnNames":            drivers.ColumnNames,
	"columnDBTypes":          drivers.ColumnDBTypes,
	"getTable":               drivers.GetTable,
}
<|endoftext|>


<commit_before>/*Package httpmw provides middlewares that are automatically
used by the generated code.*/
package httpmw

import (
	"net/http"

	"github.com/utrack/clay/transport/httptransport"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// DefaultChain is a chain that gets applied to the generated handlers.
func DefaultChain(next http.HandlerFunc) http.HandlerFunc {
	return InjectTransportStream(
		HeadersToGRPCMD(
			next,
		),
	)
}

// HeadersToGRPCMD inserts HTTP headers to gRPC metadata, as if they were
// received via gRPC.
// Every header name is lowercased, per gRPC standards.
func HeadersToGRPCMD(next http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// use metadata.FromIncomingContext to access it
		var md metadata.MD

		ctx := r.Context()
		// Use existing MD if it was injected earlier
		if m, ok := metadata.FromIncomingContext(ctx); ok {
			md = m
		} else {
			md = make(metadata.MD)
		}

		for k, v := range r.Header {
			md.Append(k, v...)
		}

		ctx = metadata.NewIncomingContext(ctx, md)

		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// InjectTransportStream injects httptransport.TransportStream to the context.
func InjectTransportStream(next http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w = httptransport.NewCodedWriter(w)

		ctx := grpc.NewContextWithServerTransportStream(r.Context(), httptransport.NewTStream(w))

		r = r.WithContext(ctx)
		next.ServeHTTP(w, r)
	})
}
<commit_msg>fixed httpransport import<commit_after>/*Package httpmw provides middlewares that are automatically
used by the generated code.*/
package httpmw

import (
	"net/http"

	"github.com/utrack/clay/v2/transport/httptransport"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// DefaultChain is a chain that gets applied to the generated handlers.
func DefaultChain(next http.HandlerFunc) http.HandlerFunc {
	return InjectTransportStream(
		HeadersToGRPCMD(
			next,
		),
	)
}

// HeadersToGRPCMD inserts HTTP headers to gRPC metadata, as if they were
// received via gRPC.
// Every header name is lowercased, per gRPC standards.
func HeadersToGRPCMD(next http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// use metadata.FromIncomingContext to access it
		var md metadata.MD

		ctx := r.Context()
		// Use existing MD if it was injected earlier
		if m, ok := metadata.FromIncomingContext(ctx); ok {
			md = m
		} else {
			md = make(metadata.MD)
		}

		for k, v := range r.Header {
			md.Append(k, v...)
		}

		ctx = metadata.NewIncomingContext(ctx, md)

		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// InjectTransportStream injects httptransport.TransportStream to the context.
func InjectTransportStream(next http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w = httptransport.NewCodedWriter(w)

		ctx := grpc.NewContextWithServerTransportStream(r.Context(), httptransport.NewTStream(w))

		r = r.WithContext(ctx)
		next.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>// +build linux

package tcp

import (
	"syscall"

	"v2ray.com/core/app/log"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet"
)

const SO_ORIGINAL_DST = 80

func GetOriginalDestination(conn internet.Connection) net.Destination {
	tcpConn, ok := conn.(internet.SysFd)
	if !ok {
		log.Trace(errors.New("Dokodemo: Failed to get sys fd."))
		return net.Destination{}
	}
	fd, err := tcpConn.SysFd()
	if err != nil {
		log.Trace(errors.New("Dokodemo: Failed to get original destination: ", err))
		return net.Destination{}
	}

	addr, err := syscall.GetsockoptIPv6Mreq(fd, syscall.IPPROTO_IP, SO_ORIGINAL_DST)
	if err != nil {
		log.Trace(errors.New("Dokodemo: Failed to call getsockopt: ", err))
		return net.Destination{}
	}
	ip := net.IPAddress(addr.Multiaddr[4:8])
	port := uint16(addr.Multiaddr[2])<<8 + uint16(addr.Multiaddr[3])
	return net.TCPDestination(ip, net.Port(port))
}
<commit_msg>fix build for linux<commit_after>// +build linux

package tcp

import (
	"syscall"

	"v2ray.com/core/app/log"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet"
)

const SO_ORIGINAL_DST = 80

func GetOriginalDestination(conn internet.Connection) net.Destination {
	tcpConn, ok := conn.(internet.SysFd)
	if !ok {
		log.Trace(errors.New("failed to get sys fd").Path("Transport", "Internet", "TCP"))
		return net.Destination{}
	}
	fd, err := tcpConn.SysFd()
	if err != nil {
		log.Trace(errors.New("failed to get original destination").Base(err).Path("Transport", "Internet", "TCP"))
		return net.Destination{}
	}

	addr, err := syscall.GetsockoptIPv6Mreq(fd, syscall.IPPROTO_IP, SO_ORIGINAL_DST)
	if err != nil {
		log.Trace(errors.New("failed to call getsockopt").Base(err).Path("Transport", "Internet", "TCP"))
		return net.Destination{}
	}
	ip := net.IPAddress(addr.Multiaddr[4:8])
	port := uint16(addr.Multiaddr[2])<<8 + uint16(addr.Multiaddr[3])
	return net.TCPDestination(ip, net.Port(port))
}
<|endoftext|>


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

package exec

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

func helperCommand(s ...string) *Cmd {
	cs := []string{"-test.run=exec.TestHelperProcess", "--"}
	cs = append(cs, s...)
	cmd := Command(os.Args[0], cs...)
	cmd.Env = append([]string{"GO_WANT_HELPER_PROCESS=1"}, os.Environ()...)
	return cmd
}

func TestEcho(t *testing.T) {
	bs, err := helperCommand("echo", "foo bar", "baz").Output()
	if err != nil {
		t.Errorf("echo: %v", err)
	}
	if g, e := string(bs), "foo bar baz\n"; g != e {
		t.Errorf("echo: want %q, got %q", e, g)
	}
}

func TestCatStdin(t *testing.T) {
	// Cat, testing stdin and stdout.
	input := "Input string\nLine 2"
	p := helperCommand("cat")
	p.Stdin = strings.NewReader(input)
	bs, err := p.Output()
	if err != nil {
		t.Errorf("cat: %v", err)
	}
	s := string(bs)
	if s != input {
		t.Errorf("cat: want %q, got %q", input, s)
	}
}

func TestCatGoodAndBadFile(t *testing.T) {
	// Testing combined output and error values.
	bs, err := helperCommand("cat", "/bogus/file.foo", "exec_test.go").CombinedOutput()
	if _, ok := err.(*ExitError); !ok {
		t.Errorf("expected *ExitError from cat combined; got %T: %v", err, err)
	}
	s := string(bs)
	sp := strings.SplitN(s, "\n", 2)
	if len(sp) != 2 {
		t.Fatalf("expected two lines from cat; got %q", s)
	}
	errLine, body := sp[0], sp[1]
	if !strings.HasPrefix(errLine, "Error: open /bogus/file.foo") {
		t.Errorf("expected stderr to complain about file; got %q", errLine)
	}
	if !strings.Contains(body, "func TestHelperProcess(t *testing.T)") {
		t.Errorf("expected test code; got %q (len %d)", body, len(body))
	}
}

func TestNoExistBinary(t *testing.T) {
	// Can't run a non-existent binary
	err := Command("/no-exist-binary").Run()
	if err == nil {
		t.Error("expected error from /no-exist-binary")
	}
}

func TestExitStatus(t *testing.T) {
	// Test that exit values are returned correctly
	err := helperCommand("exit", "42").Run()
	if werr, ok := err.(*ExitError); ok {
		if s, e := werr.Error(), "exit status 42"; s != e {
			t.Errorf("from exit 42 got exit %q, want %q", s, e)
		}
	} else {
		t.Fatalf("expected *ExitError from exit 42; got %T: %v", err, err)
	}
}

func TestPipes(t *testing.T) {
	check := func(what string, err error) {
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
	}
	// Cat, testing stdin and stdout.
	c := helperCommand("pipetest")
	stdin, err := c.StdinPipe()
	check("StdinPipe", err)
	stdout, err := c.StdoutPipe()
	check("StdoutPipe", err)
	stderr, err := c.StderrPipe()
	check("StderrPipe", err)

	outbr := bufio.NewReader(stdout)
	errbr := bufio.NewReader(stderr)
	line := func(what string, br *bufio.Reader) string {
		line, _, err := br.ReadLine()
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
		return string(line)
	}

	err = c.Start()
	check("Start", err)

	_, err = stdin.Write([]byte("O:I am output\n"))
	check("first stdin Write", err)
	if g, e := line("first output line", outbr), "O:I am output"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("E:I am error\n"))
	check("second stdin Write", err)
	if g, e := line("first error line", errbr), "E:I am error"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("O:I am output2\n"))
	check("third stdin Write 3", err)
	if g, e := line("second output line", outbr), "O:I am output2"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	stdin.Close()
	err = c.Wait()
	check("Wait", err)
}

func TestExtraFiles(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Logf("no operating system support; skipping")
		return
	}
	tf, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatalf("TempFile: %v", err)
	}
	defer os.Remove(tf.Name())
	defer tf.Close()

	const text = "Hello, fd 3!"
	_, err = tf.Write([]byte(text))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	_, err = tf.Seek(0, os.SEEK_SET)
	if err != nil {
		t.Fatalf("Seek: %v", err)
	}

	c := helperCommand("read3")
	c.ExtraFiles = []*os.File{tf}
	bs, err := c.CombinedOutput()
	if err != nil {
		t.Fatalf("CombinedOutput: %v", err)
	}
	if string(bs) != text {
		t.Errorf("got %q; want %q", string(bs), text)
	}
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "echo":
		iargs := []interface{}{}
		for _, s := range args {
			iargs = append(iargs, s)
		}
		fmt.Println(iargs...)
	case "cat":
		if len(args) == 0 {
			io.Copy(os.Stdout, os.Stdin)
			return
		}
		exit := 0
		for _, fn := range args {
			f, err := os.Open(fn)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				exit = 2
			} else {
				defer f.Close()
				io.Copy(os.Stdout, f)
			}
		}
		os.Exit(exit)
	case "pipetest":
		bufr := bufio.NewReader(os.Stdin)
		for {
			line, _, err := bufr.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				os.Exit(1)
			}
			if bytes.HasPrefix(line, []byte("O:")) {
				os.Stdout.Write(line)
				os.Stdout.Write([]byte{'\n'})
			} else if bytes.HasPrefix(line, []byte("E:")) {
				os.Stderr.Write(line)
				os.Stderr.Write([]byte{'\n'})
			} else {
				os.Exit(1)
			}
		}
	case "read3": // read fd 3
		fd3 := os.NewFile(3, "fd3")
		bs, err := ioutil.ReadAll(fd3)
		if err != nil {
			fmt.Printf("ReadAll from fd 3: %v", err)
			os.Exit(1)
		}
		os.Stderr.Write(bs)
	case "exit":
		n, _ := strconv.Atoi(args[0])
		os.Exit(n)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}
<commit_msg>os/exec: fix -test.run argument for new 'go test'<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package exec

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

func helperCommand(s ...string) *Cmd {
	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	cmd := Command(os.Args[0], cs...)
	cmd.Env = append([]string{"GO_WANT_HELPER_PROCESS=1"}, os.Environ()...)
	return cmd
}

func TestEcho(t *testing.T) {
	bs, err := helperCommand("echo", "foo bar", "baz").Output()
	if err != nil {
		t.Errorf("echo: %v", err)
	}
	if g, e := string(bs), "foo bar baz\n"; g != e {
		t.Errorf("echo: want %q, got %q", e, g)
	}
}

func TestCatStdin(t *testing.T) {
	// Cat, testing stdin and stdout.
	input := "Input string\nLine 2"
	p := helperCommand("cat")
	p.Stdin = strings.NewReader(input)
	bs, err := p.Output()
	if err != nil {
		t.Errorf("cat: %v", err)
	}
	s := string(bs)
	if s != input {
		t.Errorf("cat: want %q, got %q", input, s)
	}
}

func TestCatGoodAndBadFile(t *testing.T) {
	// Testing combined output and error values.
	bs, err := helperCommand("cat", "/bogus/file.foo", "exec_test.go").CombinedOutput()
	if _, ok := err.(*ExitError); !ok {
		t.Errorf("expected *ExitError from cat combined; got %T: %v", err, err)
	}
	s := string(bs)
	sp := strings.SplitN(s, "\n", 2)
	if len(sp) != 2 {
		t.Fatalf("expected two lines from cat; got %q", s)
	}
	errLine, body := sp[0], sp[1]
	if !strings.HasPrefix(errLine, "Error: open /bogus/file.foo") {
		t.Errorf("expected stderr to complain about file; got %q", errLine)
	}
	if !strings.Contains(body, "func TestHelperProcess(t *testing.T)") {
		t.Errorf("expected test code; got %q (len %d)", body, len(body))
	}
}

func TestNoExistBinary(t *testing.T) {
	// Can't run a non-existent binary
	err := Command("/no-exist-binary").Run()
	if err == nil {
		t.Error("expected error from /no-exist-binary")
	}
}

func TestExitStatus(t *testing.T) {
	// Test that exit values are returned correctly
	err := helperCommand("exit", "42").Run()
	if werr, ok := err.(*ExitError); ok {
		if s, e := werr.Error(), "exit status 42"; s != e {
			t.Errorf("from exit 42 got exit %q, want %q", s, e)
		}
	} else {
		t.Fatalf("expected *ExitError from exit 42; got %T: %v", err, err)
	}
}

func TestPipes(t *testing.T) {
	check := func(what string, err error) {
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
	}
	// Cat, testing stdin and stdout.
	c := helperCommand("pipetest")
	stdin, err := c.StdinPipe()
	check("StdinPipe", err)
	stdout, err := c.StdoutPipe()
	check("StdoutPipe", err)
	stderr, err := c.StderrPipe()
	check("StderrPipe", err)

	outbr := bufio.NewReader(stdout)
	errbr := bufio.NewReader(stderr)
	line := func(what string, br *bufio.Reader) string {
		line, _, err := br.ReadLine()
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
		return string(line)
	}

	err = c.Start()
	check("Start", err)

	_, err = stdin.Write([]byte("O:I am output\n"))
	check("first stdin Write", err)
	if g, e := line("first output line", outbr), "O:I am output"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("E:I am error\n"))
	check("second stdin Write", err)
	if g, e := line("first error line", errbr), "E:I am error"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("O:I am output2\n"))
	check("third stdin Write 3", err)
	if g, e := line("second output line", outbr), "O:I am output2"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	stdin.Close()
	err = c.Wait()
	check("Wait", err)
}

func TestExtraFiles(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Logf("no operating system support; skipping")
		return
	}
	tf, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatalf("TempFile: %v", err)
	}
	defer os.Remove(tf.Name())
	defer tf.Close()

	const text = "Hello, fd 3!"
	_, err = tf.Write([]byte(text))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	_, err = tf.Seek(0, os.SEEK_SET)
	if err != nil {
		t.Fatalf("Seek: %v", err)
	}

	c := helperCommand("read3")
	c.ExtraFiles = []*os.File{tf}
	bs, err := c.CombinedOutput()
	if err != nil {
		t.Fatalf("CombinedOutput: %v", err)
	}
	if string(bs) != text {
		t.Errorf("got %q; want %q", string(bs), text)
	}
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "echo":
		iargs := []interface{}{}
		for _, s := range args {
			iargs = append(iargs, s)
		}
		fmt.Println(iargs...)
	case "cat":
		if len(args) == 0 {
			io.Copy(os.Stdout, os.Stdin)
			return
		}
		exit := 0
		for _, fn := range args {
			f, err := os.Open(fn)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				exit = 2
			} else {
				defer f.Close()
				io.Copy(os.Stdout, f)
			}
		}
		os.Exit(exit)
	case "pipetest":
		bufr := bufio.NewReader(os.Stdin)
		for {
			line, _, err := bufr.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				os.Exit(1)
			}
			if bytes.HasPrefix(line, []byte("O:")) {
				os.Stdout.Write(line)
				os.Stdout.Write([]byte{'\n'})
			} else if bytes.HasPrefix(line, []byte("E:")) {
				os.Stderr.Write(line)
				os.Stderr.Write([]byte{'\n'})
			} else {
				os.Exit(1)
			}
		}
	case "read3": // read fd 3
		fd3 := os.NewFile(3, "fd3")
		bs, err := ioutil.ReadAll(fd3)
		if err != nil {
			fmt.Printf("ReadAll from fd 3: %v", err)
			os.Exit(1)
		}
		os.Stderr.Write(bs)
	case "exit":
		n, _ := strconv.Atoi(args[0])
		os.Exit(n)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}
<|endoftext|>


<commit_before>package timeutil

import (
	"errors"
	"fmt"
	"strings"
	"time"
)

/*

TimeDeltaDow is designed to retrieve a time object x days of week in the past or the future.

// Two Sundays in the future, including today, at 00:00:00
t, err := TimeDeltaDow(time.Now(), time.Sunday, 2, true, true)

// Two Sundays in the future, including today, at present time
t, err := TimeDeltaDow(time.Now(), time.Sunday, 2, true, false)

// Two Sundays ago, not including today, at 00:00:00
t, err := TimeDeltaDow(time.Now(), time.Sunday, -2, false, true)

// Two Sundays ago, not including today, at present time
t, err := TimeDeltaDow(time.Now(), time.Sunday, -2, false, false)

*/

func TimeDeltaDow(base time.Time, wantDow time.Weekday, deltaUnits int, wantInclusive bool, wantStartOfDay bool) (time.Time, error) {
	return TimeDeltaDowInt(base, int(wantDow), deltaUnits, wantInclusive, wantStartOfDay)
}

func TimeDeltaDowString(base time.Time, wantDowS string, deltaUnits int, wantInclusive bool, wantStartOfDay bool) (time.Time, error) {
	wantDow, err := ParseDayOfWeek(wantDowS)
	if err != nil {
		return time.Now(), err
	}
	return TimeDeltaDowInt(base, wantDow, deltaUnits, wantInclusive, wantStartOfDay)
}

func TimeDeltaDowInt(base time.Time, wantDow int, deltaUnits int, wantInclusive bool, wantStartOfDay bool) (time.Time, error) {
	deltaUnitsAbs := deltaUnits
	if deltaUnitsAbs < 1 {
		deltaUnitsAbs *= -1
	}
	deltaDays := int(0)
	if deltaUnits < 0 {
		deltaDaysTry, err := DaysAgoDow(int(base.Weekday()), wantDow, wantInclusive)
		if err != nil {
			return base, err
		}
		deltaDays = deltaDaysTry
	} else if deltaUnits > 0 {
		deltaDaysTry, err := DaysToDow(int(base.Weekday()), wantDow, wantInclusive)
		if err != nil {
			return base, err
		}
		deltaDays = deltaDaysTry
	} else {
		return base, errors.New("Delta units cannot be 0")
	}
	if deltaUnitsAbs > 1 {
		additional := deltaUnitsAbs - 1
		deltaDays += 7 * additional
	}
	if deltaUnits < 0 {
		deltaDays *= -1
	}
	t1 := base.AddDate(0, 0, deltaDays)
	if !wantStartOfDay {
		return t1, nil
	}
	t2 := time.Date(t1.Year(), t1.Month(), t1.Day(), 0, 0, 0, 0, t1.Location())
	return t2, nil
}

func DaysAgoDowStrings(baseDowS string, wantDowS string, wantInclusive bool) (int, error) {
	daysAgo := int(0)
	baseDow, err := ParseDayOfWeek(baseDowS)
	if err != nil {
		return daysAgo, err
	}
	wantDow, err := ParseDayOfWeek(wantDowS)
	if err != nil {
		return daysAgo, err
	}
	return DaysAgoDow(baseDow, wantDow, wantInclusive)
}

func DaysAgoDow(baseDow int, wantDow int, wantInclusive bool) (int, error) {
	if baseDow < 0 || baseDow > 6 {
		return int(0), fmt.Errorf("Base Day of Week is not in [0-6]: %v", baseDow)
	} else if wantDow < 0 || wantDow > 6 {
		return int(0), fmt.Errorf("Want Day of Week is not in [0-6]: %v", wantDow)
	}

	deltaDays1 := baseDow - wantDow
	deltaDays2 := deltaDays1
	if deltaDays2 < 0 {
		deltaDays2 += 7
	}
	if wantInclusive == false && deltaDays2 == 0 {
		deltaDays2 = 7
	}
	return deltaDays2, nil
}

func DaysToDowStrings(baseDowS string, wantDowS string, wantInclusive bool) (int, error) {
	daysAgo := int(0)
	baseDow, err := ParseDayOfWeek(baseDowS)
	if err != nil {
		return daysAgo, err
	}
	wantDow, err := ParseDayOfWeek(wantDowS)
	if err != nil {
		return daysAgo, err
	}
	return DaysToDow(baseDow, wantDow, wantInclusive)
}

func DaysToDow(baseDow int, wantDow int, wantInclusive bool) (int, error) {
	if baseDow < 0 || baseDow > 6 {
		return int(0), fmt.Errorf("Base Day of Week is not in [0-6]: %v", baseDow)
	} else if wantDow < 0 || wantDow > 6 {
		return int(0), fmt.Errorf("Want Day of Week is not in [0-6]: %v", wantDow)
	}
	deltaDays1 := wantDow - baseDow
	deltaDays2 := deltaDays1
	if deltaDays2 < 0 {
		deltaDays2 += 7
	}
	if wantInclusive == false && deltaDays2 == 0 {
		deltaDays2 = 7
	}
	return deltaDays2, nil
}

func ParseDayOfWeek(value string) (int, error) {
	valueLc := strings.ToLower(value)
	days := []string{"sunday", "monday", "tuesday", "wednesday", "thursday", "friday", "saturday"}
	mapping := map[string]int{}
	for i, dow := range days {
		mapping[dow] = int(i)
	}
	if dow, ok := mapping[valueLc]; ok {
		return dow, nil
	}
	return -1, fmt.Errorf("English name of day not found: %v", value)
}
<commit_msg>enhance timeutil.ParseDayOfWeek<commit_after>package timeutil

import (
	"errors"
	"fmt"
	"strings"
	"time"
)

/*

TimeDeltaDow is designed to retrieve a time object x days of week in the past or the future.

// Two Sundays in the future, including today, at 00:00:00
t, err := TimeDeltaDow(time.Now(), time.Sunday, 2, true, true)

// Two Sundays in the future, including today, at present time
t, err := TimeDeltaDow(time.Now(), time.Sunday, 2, true, false)

// Two Sundays ago, not including today, at 00:00:00
t, err := TimeDeltaDow(time.Now(), time.Sunday, -2, false, true)

// Two Sundays ago, not including today, at present time
t, err := TimeDeltaDow(time.Now(), time.Sunday, -2, false, false)

*/

func TimeDeltaDow(base time.Time, wantDow time.Weekday, deltaUnits int, wantInclusive bool, wantStartOfDay bool) (time.Time, error) {
	return TimeDeltaDowInt(base, int(wantDow), deltaUnits, wantInclusive, wantStartOfDay)
}

func TimeDeltaDowString(base time.Time, wantDowS string, deltaUnits int, wantInclusive bool, wantStartOfDay bool) (time.Time, error) {
	wantDow, err := ParseDayOfWeek(wantDowS)
	if err != nil {
		return time.Now(), err
	}
	return TimeDeltaDowInt(base, wantDow, deltaUnits, wantInclusive, wantStartOfDay)
}

func TimeDeltaDowInt(base time.Time, wantDow int, deltaUnits int, wantInclusive bool, wantStartOfDay bool) (time.Time, error) {
	deltaUnitsAbs := deltaUnits
	if deltaUnitsAbs < 1 {
		deltaUnitsAbs *= -1
	}
	deltaDays := int(0)
	if deltaUnits < 0 {
		deltaDaysTry, err := DaysAgoDow(int(base.Weekday()), wantDow, wantInclusive)
		if err != nil {
			return base, err
		}
		deltaDays = deltaDaysTry
	} else if deltaUnits > 0 {
		deltaDaysTry, err := DaysToDow(int(base.Weekday()), wantDow, wantInclusive)
		if err != nil {
			return base, err
		}
		deltaDays = deltaDaysTry
	} else {
		return base, errors.New("Delta units cannot be 0")
	}
	if deltaUnitsAbs > 1 {
		additional := deltaUnitsAbs - 1
		deltaDays += 7 * additional
	}
	if deltaUnits < 0 {
		deltaDays *= -1
	}
	t1 := base.AddDate(0, 0, deltaDays)
	if !wantStartOfDay {
		return t1, nil
	}
	t2 := time.Date(t1.Year(), t1.Month(), t1.Day(), 0, 0, 0, 0, t1.Location())
	return t2, nil
}

func DaysAgoDowStrings(baseDowS string, wantDowS string, wantInclusive bool) (int, error) {
	daysAgo := int(0)
	baseDow, err := ParseDayOfWeek(baseDowS)
	if err != nil {
		return daysAgo, err
	}
	wantDow, err := ParseDayOfWeek(wantDowS)
	if err != nil {
		return daysAgo, err
	}
	return DaysAgoDow(baseDow, wantDow, wantInclusive)
}

func DaysAgoDow(baseDow int, wantDow int, wantInclusive bool) (int, error) {
	if baseDow < 0 || baseDow > 6 {
		return int(0), fmt.Errorf("Base Day of Week is not in [0-6]: %v", baseDow)
	} else if wantDow < 0 || wantDow > 6 {
		return int(0), fmt.Errorf("Want Day of Week is not in [0-6]: %v", wantDow)
	}

	deltaDays1 := baseDow - wantDow
	deltaDays2 := deltaDays1
	if deltaDays2 < 0 {
		deltaDays2 += 7
	}
	if wantInclusive == false && deltaDays2 == 0 {
		deltaDays2 = 7
	}
	return deltaDays2, nil
}

func DaysToDowStrings(baseDowS string, wantDowS string, wantInclusive bool) (int, error) {
	daysAgo := int(0)
	baseDow, err := ParseDayOfWeek(baseDowS)
	if err != nil {
		return daysAgo, err
	}
	wantDow, err := ParseDayOfWeek(wantDowS)
	if err != nil {
		return daysAgo, err
	}
	return DaysToDow(baseDow, wantDow, wantInclusive)
}

func DaysToDow(baseDow int, wantDow int, wantInclusive bool) (int, error) {
	if baseDow < 0 || baseDow > 6 {
		return int(0), fmt.Errorf("Base Day of Week is not in [0-6]: %v", baseDow)
	} else if wantDow < 0 || wantDow > 6 {
		return int(0), fmt.Errorf("Want Day of Week is not in [0-6]: %v", wantDow)
	}
	deltaDays1 := wantDow - baseDow
	deltaDays2 := deltaDays1
	if deltaDays2 < 0 {
		deltaDays2 += 7
	}
	if wantInclusive == false && deltaDays2 == 0 {
		deltaDays2 = 7
	}
	return deltaDays2, nil
}

func ParseDayOfWeek(value string) (int, error) {
	valueLc := strings.ToLower(strings.TrimSpace(value))
	days := []string{"sunday", "monday", "tuesday", "wednesday", "thursday", "friday", "saturday"}
	mapping := map[string]int{}
	for i, dow := range days {
		mapping[dow] = int(i)
	}
	if dow, ok := mapping[valueLc]; ok {
		return dow, nil
	}
	return -1, fmt.Errorf("English name of day not found: %v", value)
}
<|endoftext|>


<commit_before>package nats

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func startReconnectServer(t *testing.T) *server {
	return startServer(t, 22222, "")
}

func TestReconnectDisallowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = false
	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}

	ts.stopServer()
	if e := wait(ch); e != nil {
		t.Fatal("Did not trigger ClosedCB correctly")
	}
	nc.Close()
}

func TestReconnectAllowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = true
	opts.MaxReconnect = 2
	opts.ReconnectWait = 1 * time.Second

	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ts.stopServer()

	// We want wait to timeout here, and the connection
	// should not trigger the Close CB.
	if e := wait(ch); e == nil {
		t.Fatal("Triggered ClosedCB incorrectly")
	}
	if !nc.isReconnecting() {
		t.Fatal("Expected to be in a reconnecting state")
	}

	// clear the CloseCB since ch will block
	nc.Opts.ClosedCB = nil
	nc.Close()
}

var reconnectOpts = Options{
	Url:            "nats://localhost:22222",
	AllowReconnect: true,
	MaxReconnect:   10,
	ReconnectWait:  100 * time.Millisecond,
	Timeout:        DefaultTimeout,
}

func TestBasicReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Stats.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Stats.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestExtendedReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts
	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	rch := make(chan bool)
	opts.ReconnectedCB = func(_ *Conn) {
		rch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}
	testString := "bar"
	received := int32(0)

	ec.Subscribe("foo", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	sub, _ := ec.Subscribe("foobar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	ec.Publish("foo", testString)
	ec.Flush()

	ts.stopServer()
	// server is stopped here..

	// wait for disconnect
	if e := waitTime(dch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a disconnect callback message")
	}

	// Sub while disconnected
	ec.Subscribe("bar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	// Unsub while disconnected
	sub.Unsubscribe()

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	if err = ec.Publish("bar", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	// server is restarted here..
	// wait for reconnect
	if e := waitTime(rch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a reconnect callback message")
	}

	if err = ec.Publish("foobar", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	ch := make(chan bool)
	ec.Subscribe("done", func(b bool) {
		ch <- true
	})
	ec.Publish("done", true)

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	if received != 4 {
		t.Fatalf("Received != %d, equals %d\n", 4, received)
	}
}

func TestParseStateReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	// Simulate partialState, this needs to be cleared
	nc.mu.Lock()
	nc.ps.state = OP_PON
	nc.mu.Unlock()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Stats.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Stats.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestQueueSubsOnReconnect(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts

	// Allow us to block on receonnect complete.
	reconnectsDone := make(chan bool)
	opts.ReconnectedCB = func(nc *Conn) {
		reconnectsDone <- true
	}

	// Helper to wait on a reconnect.
	waitOnReconnect := func() {
		select {
		case <-reconnectsDone:
			break
		case <-time.After(2 * time.Second):
			t.Fatalf("Expected a reconnect, timedout!\n")
		}
	}

	// Create connection
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "json")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	// To hold results.
	results := make(map[int]int)
	var mu sync.Mutex

	// Make sure we got what we needed, 1 msg only and all seqnos accounted for..
	checkResults := func(numSent int) {
		mu.Lock()
		defer mu.Unlock()

		for i := 0; i < numSent; i++ {
			if results[i] != 1 {
				t.Fatalf("Received incorrect number of messages, [%d] for seq: %d\n", results[i], i)
			}
		}

		// Auto reset results map
		results = make(map[int]int)
	}

	subj := "foo.bar"
	qgroup := "workers"

	cb := func(seqno int) {
		mu.Lock()
		defer mu.Unlock()
		results[seqno] = results[seqno] + 1
	}

	// Create Queue Subscribers
	ec.QueueSubscribe(subj, qgroup, cb)
	ec.QueueSubscribe(subj, qgroup, cb)

	ec.Flush()

	// Helper function to send messages and check results.
	sendAndCheckMsgs := func(numToSend int) {
		for i := 0; i < numToSend; i++ {
			ec.Publish(subj, i)
		}
		// Wait for processing.
		ec.Flush()
		time.Sleep(50 * time.Millisecond)

		// Check Results
		checkResults(numToSend)
	}

	// Base Test
	sendAndCheckMsgs(10)

	// Stop and restart server
	ts.stopServer()
	ts = startReconnectServer(t)
	defer ts.stopServer()

	waitOnReconnect()

	// Reconnect Base Test
	sendAndCheckMsgs(10)
}

func TestIsClosed(t *testing.T) {
	ts := startReconnectServer(t)
	nc := newConnection(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts.stopServer()
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts = startReconnectServer(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	nc.Close()
	if nc.IsClosed() == false {
		t.Fatalf("IsClosed returned false after Close() was called.")
	}
	ts.stopServer()
}
<commit_msg>spelling<commit_after>package nats

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func startReconnectServer(t *testing.T) *server {
	return startServer(t, 22222, "")
}

func TestReconnectDisallowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = false
	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}

	ts.stopServer()
	if e := wait(ch); e != nil {
		t.Fatal("Did not trigger ClosedCB correctly")
	}
	nc.Close()
}

func TestReconnectAllowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = true
	opts.MaxReconnect = 2
	opts.ReconnectWait = 1 * time.Second

	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ts.stopServer()

	// We want wait to timeout here, and the connection
	// should not trigger the Close CB.
	if e := wait(ch); e == nil {
		t.Fatal("Triggered ClosedCB incorrectly")
	}
	if !nc.isReconnecting() {
		t.Fatal("Expected to be in a reconnecting state")
	}

	// clear the CloseCB since ch will block
	nc.Opts.ClosedCB = nil
	nc.Close()
}

var reconnectOpts = Options{
	Url:            "nats://localhost:22222",
	AllowReconnect: true,
	MaxReconnect:   10,
	ReconnectWait:  100 * time.Millisecond,
	Timeout:        DefaultTimeout,
}

func TestBasicReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Stats.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Stats.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestExtendedReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts
	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	rch := make(chan bool)
	opts.ReconnectedCB = func(_ *Conn) {
		rch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}
	testString := "bar"
	received := int32(0)

	ec.Subscribe("foo", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	sub, _ := ec.Subscribe("foobar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	ec.Publish("foo", testString)
	ec.Flush()

	ts.stopServer()
	// server is stopped here..

	// wait for disconnect
	if e := waitTime(dch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a disconnect callback message")
	}

	// Sub while disconnected
	ec.Subscribe("bar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	// Unsub while disconnected
	sub.Unsubscribe()

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	if err = ec.Publish("bar", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	// server is restarted here..
	// wait for reconnect
	if e := waitTime(rch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a reconnect callback message")
	}

	if err = ec.Publish("foobar", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	ch := make(chan bool)
	ec.Subscribe("done", func(b bool) {
		ch <- true
	})
	ec.Publish("done", true)

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	if received != 4 {
		t.Fatalf("Received != %d, equals %d\n", 4, received)
	}
}

func TestParseStateReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	// Simulate partialState, this needs to be cleared
	nc.mu.Lock()
	nc.ps.state = OP_PON
	nc.mu.Unlock()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Stats.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Stats.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestQueueSubsOnReconnect(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts

	// Allow us to block on reconnect complete.
	reconnectsDone := make(chan bool)
	opts.ReconnectedCB = func(nc *Conn) {
		reconnectsDone <- true
	}

	// Helper to wait on a reconnect.
	waitOnReconnect := func() {
		select {
		case <-reconnectsDone:
			break
		case <-time.After(2 * time.Second):
			t.Fatalf("Expected a reconnect, timedout!\n")
		}
	}

	// Create connection
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "json")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	// To hold results.
	results := make(map[int]int)
	var mu sync.Mutex

	// Make sure we got what we needed, 1 msg only and all seqnos accounted for..
	checkResults := func(numSent int) {
		mu.Lock()
		defer mu.Unlock()

		for i := 0; i < numSent; i++ {
			if results[i] != 1 {
				t.Fatalf("Received incorrect number of messages, [%d] for seq: %d\n", results[i], i)
			}
		}

		// Auto reset results map
		results = make(map[int]int)
	}

	subj := "foo.bar"
	qgroup := "workers"

	cb := func(seqno int) {
		mu.Lock()
		defer mu.Unlock()
		results[seqno] = results[seqno] + 1
	}

	// Create Queue Subscribers
	ec.QueueSubscribe(subj, qgroup, cb)
	ec.QueueSubscribe(subj, qgroup, cb)

	ec.Flush()

	// Helper function to send messages and check results.
	sendAndCheckMsgs := func(numToSend int) {
		for i := 0; i < numToSend; i++ {
			ec.Publish(subj, i)
		}
		// Wait for processing.
		ec.Flush()
		time.Sleep(50 * time.Millisecond)

		// Check Results
		checkResults(numToSend)
	}

	// Base Test
	sendAndCheckMsgs(10)

	// Stop and restart server
	ts.stopServer()
	ts = startReconnectServer(t)
	defer ts.stopServer()

	waitOnReconnect()

	// Reconnect Base Test
	sendAndCheckMsgs(10)
}

func TestIsClosed(t *testing.T) {
	ts := startReconnectServer(t)
	nc := newConnection(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts.stopServer()
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts = startReconnectServer(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	nc.Close()
	if nc.IsClosed() == false {
		t.Fatalf("IsClosed returned false after Close() was called.")
	}
	ts.stopServer()
}
<|endoftext|>


<commit_before>package distributor

import (
	"errors"
	"sync"
	"time"

	"github.com/coreos/torus"
	"golang.org/x/net/context"
)

var (
	ErrNoPeersBlock = errors.New("distributor: no peers available for a block")
)

func (d *Distributor) GetBlock(ctx context.Context, i torus.BlockRef) ([]byte, error) {
	d.mut.RLock()
	defer d.mut.RUnlock()
	promDistBlockRequests.Inc()
	bcache, ok := d.readCache.Get(string(i.ToBytes()))
	if ok {
		promDistBlockCacheHits.Inc()
		return bcache.([]byte), nil
	}
	peers, err := d.ring.GetPeers(i)
	if err != nil {
		promDistBlockFailures.Inc()
		return nil, err
	}
	if len(peers.Peers) == 0 {
		promDistBlockFailures.Inc()
		return nil, ErrNoPeersBlock
	}
	writeLevel := d.getWriteFromServer()
	for _, p := range peers.Peers[:peers.Replication] {
		if p == d.UUID() || writeLevel == torus.WriteLocal {
			b, err := d.blocks.GetBlock(ctx, i)
			if err == nil {
				promDistBlockLocalHits.Inc()
				return b, nil
			}
			promDistBlockLocalFailures.Inc()
			break
		}
	}
	var blk []byte
	readLevel := d.getReadFromServer()
	switch readLevel {
	case torus.ReadBlock:
		blk, err = d.readWithBackoff(ctx, i, peers)
	case torus.ReadSequential:
		blk, err = d.readSequential(ctx, i, peers, clientTimeout)
	case torus.ReadSpread:
		blk, err = d.readSpread(ctx, i, peers)
	default:
		panic("unhandled read level")
	}
	if err != nil {
		// We completely failed!
		promDistBlockFailures.Inc()
		clog.Errorf("no peers for block %s: %v", i, err)
	}
	return blk, err
}

func (d *Distributor) readWithBackoff(ctx context.Context, ref torus.BlockRef, peers torus.PeerPermutation) ([]byte, error) {
	for i := uint(0); i < 10; i++ {
		timeout := clientTimeout * (1 << i)
		blk, err := d.readSequential(ctx, ref, peers, timeout)
		if err == nil {
			return blk, err
		}
	}
	return nil, ErrNoPeersBlock
}

func (d *Distributor) readSequential(ctx context.Context, i torus.BlockRef, peers torus.PeerPermutation, timeout time.Duration) ([]byte, error) {
	for _, p := range peers.Peers {
		// If it's local, just try to get it.
		if p == d.UUID() {
			b, err := d.blocks.GetBlock(ctx, i)
			if err == nil {
				promDistBlockLocalHits.Inc()
				return b, nil
			}
			promDistBlockLocalFailures.Inc()
			clog.Debug("failed local peer (again)")
			continue
		}
		// Fetch block from remote. First pass through peers
		// with a timeout, then through the list without.
		getctx, cancel := context.WithTimeout(ctx, timeout)
		blk, err := d.readFromPeer(getctx, i, p)
		cancel()

		if err == nil {
			return blk, nil
		}

		// If this peer didn't have it, continue
		if err == torus.ErrBlockUnavailable || err == torus.ErrNoPeer {
			clog.Warningf("block %s from %s failed, trying next peer", i, p)
			promDistBlockPeerFailures.WithLabelValues(p).Inc()
			continue
		}

		// If there was a more significant error, fail hard.
		promDistBlockFailures.Inc()
		clog.Errorf("failed remote peer %s %s %#v", p, err, err)
		return nil, err
	}
	return nil, ErrNoPeersBlock
}

func (d *Distributor) readSpread(ctx context.Context, i torus.BlockRef, peers torus.PeerPermutation) ([]byte, error) {
	resch := make(chan []byte)
	errch := make(chan error, peers.Replication)
	var once sync.Once
	count := 0
	for _, p := range peers.Peers[:peers.Replication] {
		if p == d.UUID() {
			continue
		}
		go func(peer string) {
			getctx, cancel := context.WithTimeout(ctx, clientTimeout)
			blk, err := d.readFromPeer(getctx, i, peer)
			cancel()
			if err == nil {
				once.Do(func() {
					resch <- blk
					close(resch)
				})
				return
			}
			errch <- err
		}(p)
		count++
	}

	for {
		select {
		case blk := <-resch:
			return blk, nil
		case err := <-errch:
			clog.Debugf("spread-read: %s, %s", err, i)
			count--
			if count == 0 {
				return nil, ErrNoPeersBlock
			}
		}
	}
}

func (d *Distributor) readFromPeer(ctx context.Context, i torus.BlockRef, peer string) ([]byte, error) {
	blk, err := d.client.GetBlock(ctx, peer, i)
	// If we're successful, store that.
	if err == nil {
		d.readCache.Put(string(i.ToBytes()), blk)
		promDistBlockPeerHits.WithLabelValues(peer).Inc()
		return blk, nil
	}
	return nil, err
}

func (d *Distributor) getWriteFromServer() torus.WriteLevel {
	return d.srv.Cfg.WriteLevel
}

func (d *Distributor) getReadFromServer() torus.ReadLevel {
	return d.srv.Cfg.ReadLevel
}

func (d *Distributor) WriteBlock(ctx context.Context, i torus.BlockRef, data []byte) error {
	d.mut.RLock()
	defer d.mut.RUnlock()
	peers, err := d.ring.GetPeers(i)
	if err != nil {
		return err
	}
	if len(peers.Peers) == 0 {
		return torus.ErrOutOfSpace
	}
	d.readCache.Put(string(i.ToBytes()), data)
	switch d.getWriteFromServer() {
	case torus.WriteLocal:
		err = d.blocks.WriteBlock(ctx, i, data)
		if err == nil {
			return nil
		}
		clog.Tracef("Couldn't write locally; writing to cluster")
		fallthrough
	case torus.WriteOne:
		for _, p := range peers.Peers[:peers.Replication] {
			// If we're one of the desired peers, we count, write here first.
			if p == d.UUID() {
				err = d.blocks.WriteBlock(ctx, i, data)
				if err != nil {
					clog.Noticef("WriteOne error, local: %s", err)
				} else {
					return nil
				}
			}
		}
		for _, p := range peers.Peers {
			err = d.client.PutBlock(ctx, p, i, data)
			if err == nil {
				return nil
			}
			clog.Noticef("WriteOne error, remote: %s", err)
		}
		return torus.ErrNoPeer
	case torus.WriteAll:
		toWrite := peers.Replication
		for _, p := range peers.Peers {
			var err error
			if p == d.UUID() {
				err = d.blocks.WriteBlock(ctx, i, data)
			} else {
				err = d.client.PutBlock(ctx, p, i, data)
			}
			if err != nil {
				clog.Noticef("error WriteAll to peer %s: %s", p, err)
			} else {
				toWrite--
			}
			if toWrite == 0 {
				return nil
			}
		}
		if toWrite == peers.Replication {
			clog.Noticef("error WriteAll to all peers")
			return torus.ErrNoPeer
		}
		clog.Warningf("only wrote block to %d/%d peers", (peers.Replication - toWrite), peers.Replication)
	}
	return nil
}

func (d *Distributor) WriteBuf(ctx context.Context, i torus.BlockRef) ([]byte, error) {
	return d.blocks.WriteBuf(ctx, i)
}

func (d *Distributor) HasBlock(ctx context.Context, i torus.BlockRef) (bool, error) {
	return false, errors.New("unimplemented -- finding if a block exists cluster-wide")
}

func (d *Distributor) DeleteBlock(ctx context.Context, i torus.BlockRef) error {
	return d.blocks.DeleteBlock(ctx, i)
}

func (d *Distributor) NumBlocks() uint64 {
	d.mut.RLock()
	defer d.mut.RUnlock()
	return d.blocks.NumBlocks()
}

func (d *Distributor) UsedBlocks() uint64 {
	d.mut.RLock()
	defer d.mut.RUnlock()
	return d.blocks.UsedBlocks()
}

func (d *Distributor) BlockIterator() torus.BlockIterator {
	return d.blocks.BlockIterator()
}

func (d *Distributor) Flush() error {
	return d.blocks.Flush()
}

func (d *Distributor) Kind() string { return "distributor" }

func (d *Distributor) BlockSize() uint64 {
	return d.blocks.BlockSize()
}
<commit_msg>prometheus: move promDistBlockPeerFailures to readFromPeer<commit_after>package distributor

import (
	"errors"
	"sync"
	"time"

	"github.com/coreos/torus"
	"golang.org/x/net/context"
)

var (
	ErrNoPeersBlock = errors.New("distributor: no peers available for a block")
)

func (d *Distributor) GetBlock(ctx context.Context, i torus.BlockRef) ([]byte, error) {
	d.mut.RLock()
	defer d.mut.RUnlock()
	promDistBlockRequests.Inc()
	bcache, ok := d.readCache.Get(string(i.ToBytes()))
	if ok {
		promDistBlockCacheHits.Inc()
		return bcache.([]byte), nil
	}
	peers, err := d.ring.GetPeers(i)
	if err != nil {
		promDistBlockFailures.Inc()
		return nil, err
	}
	if len(peers.Peers) == 0 {
		promDistBlockFailures.Inc()
		return nil, ErrNoPeersBlock
	}
	writeLevel := d.getWriteFromServer()
	for _, p := range peers.Peers[:peers.Replication] {
		if p == d.UUID() || writeLevel == torus.WriteLocal {
			b, err := d.blocks.GetBlock(ctx, i)
			if err == nil {
				promDistBlockLocalHits.Inc()
				return b, nil
			}
			promDistBlockLocalFailures.Inc()
			break
		}
	}
	var blk []byte
	readLevel := d.getReadFromServer()
	switch readLevel {
	case torus.ReadBlock:
		blk, err = d.readWithBackoff(ctx, i, peers)
	case torus.ReadSequential:
		blk, err = d.readSequential(ctx, i, peers, clientTimeout)
	case torus.ReadSpread:
		blk, err = d.readSpread(ctx, i, peers)
	default:
		panic("unhandled read level")
	}
	if err != nil {
		// We completely failed!
		promDistBlockFailures.Inc()
		clog.Errorf("no peers for block %s: %v", i, err)
	}
	return blk, err
}

func (d *Distributor) readWithBackoff(ctx context.Context, ref torus.BlockRef, peers torus.PeerPermutation) ([]byte, error) {
	for i := uint(0); i < 10; i++ {
		timeout := clientTimeout * (1 << i)
		blk, err := d.readSequential(ctx, ref, peers, timeout)
		if err == nil {
			return blk, err
		}
	}
	return nil, ErrNoPeersBlock
}

func (d *Distributor) readSequential(ctx context.Context, i torus.BlockRef, peers torus.PeerPermutation, timeout time.Duration) ([]byte, error) {
	for _, p := range peers.Peers {
		// If it's local, just try to get it.
		if p == d.UUID() {
			b, err := d.blocks.GetBlock(ctx, i)
			if err == nil {
				promDistBlockLocalHits.Inc()
				return b, nil
			}
			promDistBlockLocalFailures.Inc()
			clog.Debug("failed local peer (again)")
			continue
		}
		// Fetch block from remote. First pass through peers
		// with a timeout, then through the list without.
		getctx, cancel := context.WithTimeout(ctx, timeout)
		blk, err := d.readFromPeer(getctx, i, p)
		cancel()

		if err == nil {
			return blk, nil
		}

		// If this peer didn't have it, continue
		if err == torus.ErrBlockUnavailable || err == torus.ErrNoPeer {
			clog.Warningf("block %s from %s failed, trying next peer", i, p)
			continue
		}

		// If there was a more significant error, fail hard.
		promDistBlockFailures.Inc()
		clog.Errorf("failed remote peer %s %s %#v", p, err, err)
		return nil, err
	}
	return nil, ErrNoPeersBlock
}

func (d *Distributor) readSpread(ctx context.Context, i torus.BlockRef, peers torus.PeerPermutation) ([]byte, error) {
	resch := make(chan []byte)
	errch := make(chan error, peers.Replication)
	var once sync.Once
	count := 0
	for _, p := range peers.Peers[:peers.Replication] {
		if p == d.UUID() {
			continue
		}
		go func(peer string) {
			getctx, cancel := context.WithTimeout(ctx, clientTimeout)
			blk, err := d.readFromPeer(getctx, i, peer)
			cancel()
			if err == nil {
				once.Do(func() {
					resch <- blk
					close(resch)
				})
				return
			}
			errch <- err
		}(p)
		count++
	}

	for {
		select {
		case blk := <-resch:
			return blk, nil
		case err := <-errch:
			clog.Debugf("spread-read: %s, %s", err, i)
			count--
			if count == 0 {
				return nil, ErrNoPeersBlock
			}
		}
	}
}

func (d *Distributor) readFromPeer(ctx context.Context, i torus.BlockRef, peer string) ([]byte, error) {
	blk, err := d.client.GetBlock(ctx, peer, i)
	// If we're successful, store that.
	if err == nil {
		d.readCache.Put(string(i.ToBytes()), blk)
		promDistBlockPeerHits.WithLabelValues(peer).Inc()
		return blk, nil
	}
	promDistBlockPeerFailures.WithLabelValues(peer).Inc()
	return nil, err
}

func (d *Distributor) getWriteFromServer() torus.WriteLevel {
	return d.srv.Cfg.WriteLevel
}

func (d *Distributor) getReadFromServer() torus.ReadLevel {
	return d.srv.Cfg.ReadLevel
}

func (d *Distributor) WriteBlock(ctx context.Context, i torus.BlockRef, data []byte) error {
	d.mut.RLock()
	defer d.mut.RUnlock()
	peers, err := d.ring.GetPeers(i)
	if err != nil {
		return err
	}
	if len(peers.Peers) == 0 {
		return torus.ErrOutOfSpace
	}
	d.readCache.Put(string(i.ToBytes()), data)
	switch d.getWriteFromServer() {
	case torus.WriteLocal:
		err = d.blocks.WriteBlock(ctx, i, data)
		if err == nil {
			return nil
		}
		clog.Tracef("Couldn't write locally; writing to cluster")
		fallthrough
	case torus.WriteOne:
		for _, p := range peers.Peers[:peers.Replication] {
			// If we're one of the desired peers, we count, write here first.
			if p == d.UUID() {
				err = d.blocks.WriteBlock(ctx, i, data)
				if err != nil {
					clog.Noticef("WriteOne error, local: %s", err)
				} else {
					return nil
				}
			}
		}
		for _, p := range peers.Peers {
			err = d.client.PutBlock(ctx, p, i, data)
			if err == nil {
				return nil
			}
			clog.Noticef("WriteOne error, remote: %s", err)
		}
		return torus.ErrNoPeer
	case torus.WriteAll:
		toWrite := peers.Replication
		for _, p := range peers.Peers {
			var err error
			if p == d.UUID() {
				err = d.blocks.WriteBlock(ctx, i, data)
			} else {
				err = d.client.PutBlock(ctx, p, i, data)
			}
			if err != nil {
				clog.Noticef("error WriteAll to peer %s: %s", p, err)
			} else {
				toWrite--
			}
			if toWrite == 0 {
				return nil
			}
		}
		if toWrite == peers.Replication {
			clog.Noticef("error WriteAll to all peers")
			return torus.ErrNoPeer
		}
		clog.Warningf("only wrote block to %d/%d peers", (peers.Replication - toWrite), peers.Replication)
	}
	return nil
}

func (d *Distributor) WriteBuf(ctx context.Context, i torus.BlockRef) ([]byte, error) {
	return d.blocks.WriteBuf(ctx, i)
}

func (d *Distributor) HasBlock(ctx context.Context, i torus.BlockRef) (bool, error) {
	return false, errors.New("unimplemented -- finding if a block exists cluster-wide")
}

func (d *Distributor) DeleteBlock(ctx context.Context, i torus.BlockRef) error {
	return d.blocks.DeleteBlock(ctx, i)
}

func (d *Distributor) NumBlocks() uint64 {
	d.mut.RLock()
	defer d.mut.RUnlock()
	return d.blocks.NumBlocks()
}

func (d *Distributor) UsedBlocks() uint64 {
	d.mut.RLock()
	defer d.mut.RUnlock()
	return d.blocks.UsedBlocks()
}

func (d *Distributor) BlockIterator() torus.BlockIterator {
	return d.blocks.BlockIterator()
}

func (d *Distributor) Flush() error {
	return d.blocks.Flush()
}

func (d *Distributor) Kind() string { return "distributor" }

func (d *Distributor) BlockSize() uint64 {
	return d.blocks.BlockSize()
}
<|endoftext|>